{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport torch\nos.getcwd()\n\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:45:33.587548Z","iopub.execute_input":"2023-08-10T05:45:33.588130Z","iopub.status.idle":"2023-08-10T05:45:37.797866Z","shell.execute_reply.started":"2023-08-10T05:45:33.588080Z","shell.execute_reply":"2023-08-10T05:45:37.796368Z"},"trusted":true},"execution_count":1,"outputs":[]},{"cell_type":"code","source":"# 数据预处理\nproject_root = \"/kaggle/working/conveyer_belt_detector/\"\ndataset_root = \"/kaggle/input/conveyer-belt-detect/dataset/\"\n\nclass1_train_pic_root = os.path.join(dataset_root, \"class1/train\")\nclass1_val_pic_root = os.path.join(dataset_root, \"class1/val\")\nclass2_train_pic_root = os.path.join(dataset_root, \"class2/train\")\nclass2_val_pic_root = os.path.join(dataset_root, \"class2/val\")\nclass3_train_pic_root = os.path.join(dataset_root, \"class3/train\")\nclass3_val_pic_root = os.path.join(dataset_root, \"class3/val\")\nnormal_train_pic_root = os.path.join(dataset_root, \"normal/train\")\nnormal_val_pic_root = os.path.join(dataset_root, \"normal/val\")\n\n# 任务一的标注文件\ntask1_train_class1_orig_anno = os.path.join(dataset_root, \"class1/train_infos.csv\")\ntask1_train_class2_orig_anno = os.path.join(dataset_root, \"class2/train_infos.csv\")\ntask1_train_class3_orig_anno = os.path.join(dataset_root, \"class3/train_infos.csv\")\ntask1_train_normal_orig_anno = os.path.join(dataset_root, \"normal/train_infos.csv\")\n\ntask1_val_class1_orig_anno = os.path.join(dataset_root, \"class1/val_infos.csv\")\ntask1_val_class2_orig_anno = os.path.join(dataset_root, \"class2/val_infos.csv\")\ntask1_val_class3_orig_anno = os.path.join(dataset_root, \"class3/val_infos.csv\")\ntask1_val_normal_orig_anno = os.path.join(dataset_root, \"normal/val_infos.csv\")\n\n# 任务二的标注文件\ntask2_train_class1_orig_anno = os.path.join(dataset_root, \"class1/class1_train.txt\")\ntask2_train_class2_orig_anno = os.path.join(dataset_root, \"class2/class2_train.txt\")\ntask2_train_class3_orig_anno = os.path.join(project_root, \"tmp/class3/class3_train.txt\")\ntask2_train_normal_orig_anno = os.path.join(dataset_root, \"normal/normal_train.txt\")\n\ntask2_val_class1_orig_anno = os.path.join(dataset_root, \"class1/class1_val.txt\")\ntask2_val_class2_orig_anno = os.path.join(dataset_root, \"class2/clsss2_val.txt\")\ntask2_val_class3_orig_anno = os.path.join(project_root, \"tmp/class3/class3_val.txt\")\ntask2_val_normal_orig_anno = os.path.join(dataset_root, \"normal/normal_val.txt\")\n\n# 186 2126 2507\n\n# 汇总的标注文件\ntask1_train_detect_anno = \"/kaggle/input/anno-data-new/train_info_all.csv\"\ntask1_val_detect_anno = os.path.join(project_root, \"data/task1/val/val_infos.csv\")\ntask2_train_classes_anno = \"/kaggle/input/anno-data-new/classes_all.txt\"\ntask2_val_classes_anno = os.path.join(project_root, \"data/task2/val/classes.txt\")\n\npretrain_path=\"/kaggle/input/nvidia-ssdpyt-fp32-190826pt/nvidia_ssdpyt_fp32_190826.pt\"\n\ndef make_dir(path):\n    dir_name = os.path.dirname(path)\n    if not os.path.exists(dir_name):\n        print(\"创建目录\", dir_name)\n        os.makedirs(dir_name)\n\nmake_dir(path=project_root)\nmake_dir(path=task1_train_detect_anno)\nmake_dir(path=task1_val_detect_anno)\nmake_dir(path=task2_train_classes_anno)\nmake_dir(path=task2_val_classes_anno)\n\n# 预处理任务一的数据：\n# 1、将train_infos.csv的filename字段改为绝对路径;\n# 2、将xmin、xmax、ymin、ymax的坐标值改为相对坐标比例;\n# 3、将class1、class2的train_infos.csv合并为一个大的train_infos.csv。\nimport csv\nfrom collections import defaultdict\n\ndef predeal_task1_anno_file(dest_file, orig_files, pic_roots):\n    dest_file_dir = os.path.dirname(p=dest_file)\n    if not os.path.exists(dest_file_dir):\n        print(f\"创建目录{dest_file_dir}\")\n        os.makedirs(dest_file_dir)\n    \n    with open(file=dest_file, mode=\"w\", encoding=\"utf8\", newline='') as dest_csv_file:\n        header = None\n        csv_writer = None\n        last_file_index = 0\n        cur_index = 0\n        for orig_file, pic_root in zip(orig_files, pic_roots):\n            with open(file=orig_file, mode=\"r\", encoding=\"utf8\") as tmp_csv_file:\n                # 读表头\n                raw_line = tmp_csv_file.readline()\n                if len(raw_line) <= 0:\n                    # 读取下个文件\n                    continue\n\n                if header is None:\n                    header = raw_line.strip().split(\",\")\n                    csv_writer = csv.DictWriter(f=dest_csv_file, fieldnames=header)\n                    csv_writer.writeheader()\n                        \n                while True:\n                    raw_line = tmp_csv_file.readline()\n                    if len(raw_line) <= 0:\n                        last_file_index = cur_index + 1\n                        # 读取下个文件\n                        break\n                    line_content = raw_line.strip().split(\",\")\n\n                    new_row = defaultdict()\n                    for i in range(len(header)):\n                        if header[i] == \"index\":\n                            cur_index = (int)(line_content[0])\n                            new_row[header[i]] = last_file_index + cur_index\n                        elif header[i] == \"filename\":\n                            new_row[header[i]] = os.path.join(pic_root, line_content[i])\n                        elif header[i] == \"xmin\" or header[i] == \"xmax\":\n                            new_row[header[i]] = float(line_content[i]) / float(line_content[8])\n                        elif header[i] == \"ymin\" or header[i] == \"ymax\":\n                            new_row[header[i]] = float(line_content[i]) / float(line_content[7])\n                        else:\n                            new_row[header[i]] = line_content[i]\n\n\n                    csv_writer.writerow(new_row)\n        print(\"导出 \" + dest_file + \" 成功\")\n\n# 预处理class1、class2、class3、normal的标注文件\ntrain_anno_files = [task1_train_class1_orig_anno, task1_train_class2_orig_anno]\ntrain_pic_roots = [class1_train_pic_root, class2_train_pic_root] \n\nval_anno_files = [task1_val_class1_orig_anno, task1_val_class2_orig_anno]\nval_pic_roots = [class1_val_pic_root, class2_val_pic_root] \n\n# predeal_task1_anno_file(dest_file=task1_train_detect_anno, orig_files=train_anno_files, pic_roots=train_pic_roots)\n# predeal_task1_anno_file(dest_file=task1_val_detect_anno, orig_files=val_anno_files, pic_roots=val_pic_roots)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:45:37.800837Z","iopub.execute_input":"2023-08-10T05:45:37.802020Z","iopub.status.idle":"2023-08-10T05:45:37.837414Z","shell.execute_reply.started":"2023-08-10T05:45:37.801964Z","shell.execute_reply":"2023-08-10T05:45:37.836329Z"},"trusted":true},"execution_count":2,"outputs":[]},{"cell_type":"code","source":"# 预处理任务二的数据\ndef generate_task2_class3_anno(dest_files, orig_dirs, append_dirs):\n    header = [\"path\", \"class\"]\n    for dest_file, orig_dir, append_dir in zip(dest_files, orig_dirs, append_dirs):\n        make_dir(dest_file)\n        with open(file=dest_file, mode=\"w\", encoding=\"utf8\", newline=\"\") as dest_file_handle:\n            csv_writer = csv.DictWriter(f=dest_file_handle, fieldnames=header)\n            for file in os.listdir(orig_dir):\n                new_row = defaultdict()\n                new_row[\"path\"] = os.path.join(append_dir, file)\n                new_row[\"class\"] = 3\n                csv_writer.writerow(new_row)\n        print(f\"导出{dest_file}成功\")\n                \n\ndef predeal_task2_anno_file(dest_file, orig_files, dataset_root):\n    dest_file_dir = os.path.dirname(p=dest_file)\n    if not os.path.exists(dest_file_dir):\n        print(f\"创建目录{dest_file_dir}\")\n        os.makedirs(dest_file_dir)\n    \n    with open(file=dest_file, mode=\"w\", encoding=\"utf8\", newline='') as dest_csv_file:\n        header = [\"path\", \"class\"]\n        csv_writer = csv.DictWriter(f=dest_csv_file, fieldnames=header)\n        for orig_file in orig_files:\n            with open(file=orig_file, mode=\"r\", encoding=\"utf8\") as tmp_csv_file:\n                while True:\n                    raw_line = tmp_csv_file.readline()\n                    if len(raw_line) <= 0:\n                        # 读取下个文件\n                        break\n                    line_content = raw_line.strip().split(\",\")\n\n                    new_row = defaultdict()\n                    new_row[\"path\"] = os.path.join(dataset_root, line_content[0][2:])\n                    new_row[\"class\"] = int(line_content[1])\n                    \n                    csv_writer.writerow(new_row)\n        print(\"导出 \" + dest_file + \" 成功\")\n\ndest_files = [task2_train_class3_orig_anno, task2_val_class3_orig_anno]\norig_dirs = [class3_train_pic_root, class3_val_pic_root]\nappend_dirs = [\"./class3/train\", \"./class3/val\"]\n# 生成任务二分类3的标注文件\n# generate_task2_class3_anno(dest_files=dest_files, orig_dirs=orig_dirs, append_dirs=append_dirs)\n        \n# 合并class1、class2、class3、normal的标注文件\n# train_orig_files = [task2_train_class1_orig_anno, task2_train_class2_orig_anno, task2_train_class3_orig_anno, task2_train_normal_orig_anno] \n# val_orig_files = [task2_val_class1_orig_anno, task2_val_class2_orig_anno, task2_val_class3_orig_anno, task2_val_normal_orig_anno] \n\n# predeal_task2_anno_file(dest_file=task2_train_classes_anno, orig_files=train_orig_files, dataset_root=dataset_root)\n# predeal_task2_anno_file(dest_file=task2_val_classes_anno, orig_files=val_orig_files, dataset_root=dataset_root)\n","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:45:37.839759Z","iopub.execute_input":"2023-08-10T05:45:37.840238Z","iopub.status.idle":"2023-08-10T05:45:37.859475Z","shell.execute_reply.started":"2023-08-10T05:45:37.840203Z","shell.execute_reply":"2023-08-10T05:45:37.858219Z"},"trusted":true},"execution_count":3,"outputs":[]},{"cell_type":"code","source":"!pip install pycocotools","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:45:37.862539Z","iopub.execute_input":"2023-08-10T05:45:37.862968Z","iopub.status.idle":"2023-08-10T05:46:26.323272Z","shell.execute_reply.started":"2023-08-10T05:45:37.862901Z","shell.execute_reply":"2023-08-10T05:46:26.321765Z"},"trusted":true},"execution_count":4,"outputs":[{"name":"stdout","text":"Collecting pycocotools\n  Downloading pycocotools-2.0.6.tar.gz (24 kB)\n  Installing build dependencies ... \u001b[?25ldone\n\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n\u001b[?25hRequirement already satisfied: matplotlib>=2.1.0 in /opt/conda/lib/python3.10/site-packages (from pycocotools) (3.7.1)\nRequirement already satisfied: numpy in /opt/conda/lib/python3.10/site-packages (from pycocotools) (1.23.5)\nRequirement already satisfied: contourpy>=1.0.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.1.0->pycocotools) (1.1.0)\nRequirement already satisfied: cycler>=0.10 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.1.0->pycocotools) (0.11.0)\nRequirement already satisfied: fonttools>=4.22.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.1.0->pycocotools) (4.40.0)\nRequirement already satisfied: kiwisolver>=1.0.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.1.0->pycocotools) (1.4.4)\nRequirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.1.0->pycocotools) (21.3)\nRequirement already satisfied: pillow>=6.2.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.1.0->pycocotools) (9.5.0)\nRequirement already satisfied: pyparsing>=2.3.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.1.0->pycocotools) (3.0.9)\nRequirement already satisfied: python-dateutil>=2.7 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.1.0->pycocotools) (2.8.2)\nRequirement already satisfied: six>=1.5 in /opt/conda/lib/python3.10/site-packages (from python-dateutil>=2.7->matplotlib>=2.1.0->pycocotools) (1.16.0)\nBuilding wheels for collected packages: pycocotools\n  Building wheel for pycocotools (pyproject.toml) ... \u001b[?25ldone\n\u001b[?25h  Created wheel for pycocotools: filename=pycocotools-2.0.6-cp310-cp310-linux_x86_64.whl size=110796 sha256=019078fda887fd000777b96f70d3e30bc40739e131522b08976c111849002223\n  Stored in directory: /root/.cache/pip/wheels/58/e6/f9/f87c8f8be098b51b616871315318329cae12cdb618f4caac93\nSuccessfully built pycocotools\nInstalling collected packages: pycocotools\nSuccessfully installed pycocotools-2.0.6\n","output_type":"stream"}]},{"cell_type":"code","source":"from torchvision.transforms import ToPILImage\n\nfrom PIL import Image\nfrom PIL.Image import fromarray\nimport PIL.ImageDraw as ImageDraw\nimport PIL.ImageFont as ImageFont\nfrom PIL import ImageColor\n\nimport datetime\nimport matplotlib.pyplot as plt\n\nimport numpy as np\n\nSTANDARD_COLORS = [\n    'White', 'Cyan', 'Yellow', 'AliceBlue', 'Chartreuse', 'Aqua', 'Aquamarine', 'Azure',\n    'BlanchedAlmond', 'BlueViolet', 'BurlyWood', 'CadetBlue', 'AntiqueWhite',\n    'Chocolate', 'Coral', 'CornflowerBlue', 'Cornsilk', 'Crimson',\n    'DarkCyan', 'DarkGoldenRod', 'DarkGrey', 'DarkKhaki', 'DarkOrange',\n    'DarkOrchid', 'DarkSalmon', 'DarkSeaGreen', 'DarkTurquoise', 'DarkViolet',\n    'DeepPink', 'DeepSkyBlue', 'DodgerBlue', 'FireBrick', 'FloralWhite',\n    'ForestGreen', 'Fuchsia', 'Gainsboro', 'GhostWhite', 'Gold', 'GoldenRod',\n    'Salmon', 'Tan', 'HoneyDew', 'HotPink', 'IndianRed', 'Ivory', 'Khaki',\n    'Lavender', 'LavenderBlush', 'LawnGreen', 'LemonChiffon', 'LightBlue',\n    'LightCoral', 'LightCyan', 'LightGoldenRodYellow', 'LightGray', 'LightGrey',\n    'LightGreen', 'LightPink', 'LightSalmon', 'LightSeaGreen', 'LightSkyBlue',\n    'LightSlateGray', 'LightSlateGrey', 'LightSteelBlue', 'LightYellow', 'Lime',\n    'LimeGreen', 'Linen', 'Magenta', 'MediumAquaMarine', 'MediumOrchid',\n    'MediumPurple', 'MediumSeaGreen', 'MediumSlateBlue', 'MediumSpringGreen',\n    'MediumTurquoise', 'MediumVioletRed', 'MintCream', 'MistyRose', 'Moccasin',\n    'NavajoWhite', 'OldLace', 'Olive', 'OliveDrab', 'Orange', 'OrangeRed',\n    'Orchid', 'PaleGoldenRod', 'PaleGreen', 'PaleTurquoise', 'PaleVioletRed',\n    'PapayaWhip', 'PeachPuff', 'Peru', 'Pink', 'Plum', 'PowderBlue', 'Purple',\n    'Red', 'RosyBrown', 'RoyalBlue', 'SaddleBrown', 'SandyBrown',\n    'SeaGreen', 'SeaShell', 'Sienna', 'Silver', 'SkyBlue', 'SlateBlue',\n    'SlateGray', 'SlateGrey', 'Snow', 'SpringGreen', 'SteelBlue', 'GreenYellow',\n    'Teal', 'Thistle', 'Tomato', 'Turquoise', 'Wheat', 'WhiteSmoke', 'YellowGreen',\n    'Beige', 'Bisque', 'Violet', 'Green', 'Red'\n]\n\ndef draw_text(draw,\n              box: list,\n              cls: int,\n              score: float,\n              category_index: dict,\n              color: str,\n              font: str = 'arial.ttf',\n              font_size: int = 24):\n    \"\"\"\n    将目标边界框和类别信息绘制到图片上\n    \"\"\"\n    try:\n        font = ImageFont.truetype(font, font_size)\n    except IOError:\n        font = ImageFont.load_default()\n\n    left, top, right, bottom = box\n    # If the total height of the display strings added to the top of the bounding\n    # box exceeds the top of the image, stack the strings below the bounding box\n    # instead of above.\n    display_str = f\"{category_index[str(cls)]}: {int(100 * score)}%\"\n    display_str_heights = [font.getsize(ds)[1] for ds in display_str]\n    # Each display_str has a top and bottom margin of 0.05x.\n    display_str_height = (1 + 2 * 0.05) * max(display_str_heights)\n\n    if top > display_str_height:\n        text_top = top - display_str_height\n        text_bottom = top\n    else:\n        text_top = bottom\n        text_bottom = bottom + display_str_height\n\n    for ds in display_str:\n        text_width, text_height = font.getsize(ds)\n        margin = np.ceil(0.05 * text_width)\n        draw.rectangle([(left, text_top),\n                        (left + text_width + 2 * margin, text_bottom)], fill=color)\n        draw.text((left + margin, text_top),\n                  ds,\n                  fill='black',\n                  font=font)\n        left += text_width\n\n\ndef draw_masks(image, masks, colors, thresh: float = 0.7, alpha: float = 0.5):\n    np_image = np.array(image)\n    masks = np.where(masks > thresh, True, False)\n\n    # colors = np.array(colors)\n    img_to_draw = np.copy(np_image)\n    # TODO: There might be a way to vectorize this\n    for mask, color in zip(masks, colors):\n        img_to_draw[mask] = color\n\n    out = np_image * (1 - alpha) + img_to_draw * alpha\n    return fromarray(out.astype(np.uint8))\n\ndef draw_objs(image: Image,\n              boxes: np.ndarray = None,\n              classes: np.ndarray = None,\n              scores: np.ndarray = None,\n              masks: np.ndarray = None,\n              category_index: dict = None,\n              box_thresh: float = 0.1,\n              mask_thresh: float = 0.5,\n              line_thickness: int = 8,\n              font: str = 'arial.ttf',\n              font_size: int = 24,\n              draw_boxes_on_image: bool = True,\n              draw_masks_on_image: bool = False):\n    \"\"\"\n    将目标边界框信息，类别信息，mask信息绘制在图片上\n    Args:\n        image: 需要绘制的图片\n        boxes: 目标边界框信息\n        classes: 目标类别信息\n        scores: 目标概率信息\n        masks: 目标mask信息\n        category_index: 类别与名称字典\n        box_thresh: 过滤的概率阈值\n        mask_thresh:\n        line_thickness: 边界框宽度\n        font: 字体类型\n        font_size: 字体大小\n        draw_[boxes]_on_image:\n        draw_masks_on_image:\n\n    Returns:\n\n    \"\"\"\n\n    # 过滤掉低概率的目标\n    idxs = np.greater(scores, box_thresh)\n    \n    boxes = boxes[idxs]\n    classes = classes[idxs]\n    scores = scores[idxs]\n    if masks is not None:\n        masks = masks[idxs]\n    if len(boxes) == 0:\n        return image\n\n    colors = [ImageColor.getrgb(STANDARD_COLORS[cls % len(STANDARD_COLORS)]) for cls in classes]\n\n    if draw_boxes_on_image:\n        draw_boxes = []\n        for box, cls, score, color in zip(boxes, classes, scores, colors):\n            draw_boxes.append((box, cls, score, color))\n\n        # Draw all boxes onto image.\n        draw = ImageDraw.Draw(image)\n        for draw_box in draw_boxes[-1::-1]:\n            left, top, right, bottom = draw_box[0]\n            # 绘制目标边界框\n            draw.line([(left, top), (left, bottom), (right, bottom),\n                       (right, top), (left, top)], width=line_thickness, fill=draw_box[3])\n            \n            if category_index is not None:\n                # 绘制类别和概率信息\n                draw_text(draw, draw_box[0].tolist(), int(draw_box[1]), float(draw_box[2]), category_index, draw_box[3], font, font_size)\n\n    if draw_masks_on_image and (masks is not None):\n        # Draw all mask onto image.\n        image = draw_masks(image, masks, colors, mask_thresh)\n\n    return image\n\ndef draw_image(img_file, predict_boxes, predict_classes, predict_scores, category_index,\n                show_image_flag=True, save_image_flag=False, save_image_dir = None):\n    original_img = Image.open(img_file)\n    \n    plot_img = draw_objs(original_img,\n                         predict_boxes,\n                         predict_classes,\n                         predict_scores,\n                         category_index=category_index,\n                         box_thresh=0.5,\n                         line_thickness=2,\n                         font='simsun.ttc',\n                         font_size=20)\n    if show_image_flag:\n        plt.imshow(plot_img)\n        plt.show()\n        \n    if save_image_flag:\n        make_dir(path=save_image_dir)\n        save_file = os.path.join(save_image_dir, os.path.basename(img_file))\n        \n        plot_img.save(save_file)\n\ndef plot_loss_and_lr(train_loss, learning_rate):\n    try:\n        x = list(range(len(train_loss)))\n        fig, ax1 = plt.subplots(1, 1)\n        ax1.plot(x, train_loss, 'r', label='loss')\n        ax1.set_xlabel(\"epoch\")\n        ax1.set_ylabel(\"loss\")\n        ax1.set_title(\"Train Loss and lr\")\n        plt.legend(loc='best')\n\n        ax2 = ax1.twinx()\n        ax2.plot(x, learning_rate, label='lr')\n        ax2.set_ylabel(\"learning rate\")\n        ax2.set_xlim(0, len(train_loss))  # 设置横坐标整数间隔\n        plt.legend(loc='best')\n\n        handles1, labels1 = ax1.get_legend_handles_labels()\n        handles2, labels2 = ax2.get_legend_handles_labels()\n        plt.legend(handles1 + handles2, labels1 + labels2, loc='upper right')\n\n        fig.subplots_adjust(right=0.8)  # 防止出现保存图片显示不全的情况\n        fig.savefig('./loss_and_lr{}.png'.format(datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")))\n        plt.close()\n        print(\"successful save loss curve! \")\n    except Exception as e:\n        print(e)\n\n\ndef plot_map(mAP):\n    try:\n        x = list(range(len(mAP)))\n        plt.plot(x, mAP, label='mAp')\n        plt.xlabel('epoch')\n        plt.yticks([0.00, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40, 0.45, 0.50,\n                    0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85, 0.90, 0.95, 1.00])\n        plt.ylabel('mAP')\n        plt.title('Eval mAP')\n        plt.xlim(0, len(mAP))\n        plt.legend(loc='best')\n        plt.savefig('./mAP.png')\n        plt.close()\n        print(\"successful save mAP curve!\")\n    except Exception as e:\n        print(e)\n\ndef save_pic(image, path, bboxes=None, labels=None, height=None, width=None, category_index=None, bboxes_scores=None):\n    make_dir(path)\n    image.save(path)\n    \n    if bboxes is not None and labels is not None:\n        bboxes_out = bboxes.copy()\n        \n        bboxes_out[:, 0], bboxes_out[:, 2] =  bboxes[:, 0] * width, bboxes[:, 2] * width\n        bboxes_out[:, 1], bboxes_out[:, 3] =  bboxes[:, 1] * height, bboxes[:, 3] * height\n\n        gt_scores = np.array([1.0 for i in range(len(labels))]) if bboxes_scores is None else bboxes_scores\n\n        draw_image(path, bboxes_out, labels, gt_scores, category_index, save_image_flag=True, save_image_dir=os.path.dirname(path))\n    \ndef save_anchor_pic(image, path, bboxes, height, width, targets=None, targets_label=None, choice=None, bboxes_label=None, bboxes_scores=None):\n    toPilImage = ToPILImage()\n    if choice is None:\n        choice = 0\n    image = toPilImage(image[choice, :, :, :])\n    \n    labels = np.array([0 for _ in range(len(bboxes))]) if bboxes_label is None else bboxes_label\n    \n    if targets is not None:\n        bboxes = torch.cat((targets, bboxes.to(device=device)), dim=0).to(device=device)\n        color_count = len(STANDARD_COLORS)\n        targets_label = [color_count - 1 for _ in range(len(targets))] if targets_label is None else color_count - targets_label\n\n        labels = np.append(targets_label.cpu(), labels.cpu())\n        \n    if bboxes_scores is not None:\n        scores = []\n        \n        bboxes_scores = bboxes_scores.tolist()\n        target_scores = [1 for _ in range(len(targets))]\n        \n        scores.extend(target_scores)\n        scores.extend(bboxes_scores)\n        \n        bboxes_scores = np.array(scores)\n    \n    save_pic(image=image, path=path, bboxes=bboxes.cpu().numpy(), labels=labels, height=height, width=width, bboxes_scores=bboxes_scores)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:46:26.325715Z","iopub.execute_input":"2023-08-10T05:46:26.326125Z","iopub.status.idle":"2023-08-10T05:46:26.752377Z","shell.execute_reply.started":"2023-08-10T05:46:26.326086Z","shell.execute_reply":"2023-08-10T05:46:26.751212Z"},"trusted":true},"execution_count":5,"outputs":[]},{"cell_type":"code","source":"# src/res50_backbone.py\n\nimport torch.nn as nn\nimport torch\n\nclass Bottleneck(nn.Module):\n    expansion = 4\n\n    def __init__(self, in_channel, out_channel, stride=1, downsample=None):\n        super(Bottleneck, self).__init__()\n        self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel,\n                               kernel_size=1, stride=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(out_channel)\n\n        # -----------------------------------------\n\n        self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel,\n                               kernel_size=3, stride=stride, bias=False, padding=1)\n        self.bn2 = nn.BatchNorm2d(out_channel)\n\n        # -----------------------------------------\n\n        self.conv3 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel * self.expansion,\n                               kernel_size=1, stride=1, bias=False)\n        self.bn3 = nn.BatchNorm2d(out_channel * self.expansion)\n        self.relu = nn.ReLU(inplace=True)\n        self.downsample = downsample\n\n    def forward(self, x):\n        identity = x\n        if self.downsample is not None:\n            identity = self.downsample(x)\n\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n\n        out = self.conv2(out)\n        out = self.bn2(out)\n        out = self.relu(out)\n\n        out = self.conv3(out)\n        out = self.bn3(out)\n\n        out += identity\n        out = self.relu(out)\n\n        return out\n\n\nclass ResNet(nn.Module):\n\n    def __init__(self, block, blocks_num, num_classes=1000, include_top=True):\n        super(ResNet, self).__init__()\n        self.include_top = include_top\n        self.in_channel = 64\n        self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2,\n                               padding=3, bias=False)\n        self.bn1 = nn.BatchNorm2d(self.in_channel)\n        self.relu = nn.ReLU(inplace=True)\n        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n\n        self.layer1 = self._make_layer(block, 64, blocks_num[0])\n        self.layer2 = self._make_layer(block, 128, blocks_num[1], stride=2)\n        self.layer3 = self._make_layer(block, 256, blocks_num[2], stride=2)\n        self.layer4 = self._make_layer(block, 512, blocks_num[3], stride=2)\n\n        # 输出层\n        if self.include_top:\n            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n            self.fc = nn.Linear(512 * block.expansion, num_classes)\n\n        # 参数初始化\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n\n    def _make_layer(self, block, channel, block_num, stride=1):\n        \"\"\"\n            构建模块\n        \"\"\"\n        downsample = None\n\n        if stride != 1 or self.in_channel != channel * block.expansion:\n            downsample = nn.Sequential(\n                nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False),\n                nn.BatchNorm2d(channel * block.expansion))\n\n        layers = []\n        layers.append(block(self.in_channel, channel, downsample=downsample, stride=stride))\n        self.in_channel = channel * block.expansion\n\n        for _ in range(1, block_num):\n            layers.append(block(self.in_channel, channel))\n\n        return nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.bn1(x)\n        x = self.relu(x)\n        x = self.maxpool(x)\n\n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n        x = self.layer4(x)\n\n        # 最终输出\n        if self.include_top:\n            x = self.avgpool(x)\n            x = torch.flatten(x, 1)\n            x = self.fc(x)\n\n        return x\n\n\ndef resnet50(num_classes=1000, include_top=True):\n    return ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, include_top=include_top)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:46:26.754327Z","iopub.execute_input":"2023-08-10T05:46:26.755452Z","iopub.status.idle":"2023-08-10T05:46:26.784665Z","shell.execute_reply.started":"2023-08-10T05:46:26.755410Z","shell.execute_reply":"2023-08-10T05:46:26.783605Z"},"trusted":true},"execution_count":6,"outputs":[]},{"cell_type":"code","source":"# src/utils.py\n\nfrom math import sqrt\nimport itertools\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.jit.annotations import Tuple, List\nfrom torch import nn, Tensor\nimport numpy as np\n\ndef box_area(boxes):\n    \"\"\"\n    Computes the area of a set of bounding boxes, which are specified by its\n    (x1, y1, x2, y2) coordinates.\n\n    Arguments:\n        boxes (Tensor[N, 4]): boxes for which the area will be computed. They\n            are expected to be in (x1, y1, x2, y2) format\n\n    Returns:\n        area (Tensor[N]): area for each box\n    \"\"\"\n    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])\n\n\ndef calc_iou_tensor(boxes1, boxes2):\n    \"\"\"\n    Return intersection-over-union (Jaccard index) of boxes.\n\n    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.\n\n    Arguments:\n        boxes1 (Tensor[N, 4])\n        boxes2 (Tensor[M, 4])\n\n    Returns:\n        iou (Tensor[N, M]): the NxM matrix containing the pairwise\n            IoU values for every element in boxes1 and boxes2\n    \"\"\"\n    area1 = box_area(boxes1)\n    area2 = box_area(boxes2)\n\n    #  When the shapes do not match,\n    #  the shape of the returned output tensor follows the broadcasting rules\n    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # left-top [N,M,2]\n    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # right-bottom [N,M,2]\n\n    wh = (rb - lt).clamp(min=0)  # [N,M,2]\n    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]\n\n    iou = inter / (area1[:, None] + area2 - inter)\n    return iou\n\n# This function is from https://github.com/kuangliu/pytorch-ssd.\nclass Encoder(object):\n    \"\"\"\n        Inspired by https://github.com/kuangliu/pytorch-src\n        Transform between (bboxes, lables) <-> SSD output\n\n        dboxes: default boxes in size 8732 x 4,\n            encoder: input ltrb format, output xywh format\n            decoder: input xywh format, output ltrb format\n\n        encode:\n            input  : bboxes_in (Tensor nboxes x 4), labels_in (Tensor nboxes)\n            output : bboxes_out (Tensor 8732 x 4), labels_out (Tensor 8732)\n            criteria : IoU threshold of bboexes\n\n        decode:\n            input  : bboxes_in (Tensor 8732 x 4), scores_in (Tensor 8732 x nitems)\n            output : bboxes_out (Tensor nboxes x 4), labels_out (Tensor nboxes)\n            criteria : IoU threshold of bboexes\n            max_output : maximum number of output bboxes\n    \"\"\"\n    def __init__(self, dboxes):\n        self.dboxes = dboxes(order='ltrb')\n        self.dboxes_xywh = dboxes(order='xywh').unsqueeze(dim=0)\n        self.nboxes = self.dboxes.size(0)  # default boxes的数量\n        self.scale_xy = dboxes.scale_xy\n        self.scale_wh = dboxes.scale_wh\n\n    def encode(self, bboxes_in, labels_in, criteria=0.5):\n        \"\"\"\n        encode:\n            input  : bboxes_in (Tensor nboxes x 4), labels_in (Tensor nboxes)\n            output : bboxes_out (Tensor 8732 x 4), labels_out (Tensor 8732)\n            criteria : IoU threshold of bboexes\n        \"\"\"\n        # [nboxes, 8732]\n        bboxes_in = bboxes_in.to(device=device)\n        dboxes = self.dboxes.to(device=device)\n        ious = calc_iou_tensor(bboxes_in, dboxes)  # 计算每个GT与default box的iou\n        # [8732,]    每个锚框最可能的类型\n        best_dbox_ious, best_dbox_idx = ious.max(dim=0)  # 寻找每个default box匹配到的最大IoU\n        # [nboxes,]  每个类型对应的最佳锚框\n        best_bbox_ious, best_bbox_idx = ious.max(dim=1)  # 寻找每个GT匹配到的最大IoU\n\n        # 将每个GT匹配到的最佳default box设置为正样本（对应论文中Matching strategy的第一条）\n        # set best ious 2.0\n        best_dbox_ious.index_fill_(0, best_bbox_idx, 2.0)  # dim, index, value\n        # 将相应default box匹配最大IOU的GT索引进行替换\n        idx = torch.arange(0, best_bbox_idx.size(0), dtype=torch.int64).to(device=device)\n        best_dbox_idx[best_bbox_idx[idx]] = idx\n\n        # filter IoU > 0.5\n        # 寻找与GT iou大于0.5的default box,对应论文中Matching strategy的第二条(这里包括了第一条匹配到的信息)\n        masks = best_dbox_ious > criteria\n        # [8732,]\n        labels_out = torch.zeros(self.nboxes, dtype=torch.int64).to(device=device)\n        labels_out[masks] = labels_in[best_dbox_idx[masks]]\n        # 将default box匹配到正样本的位置设置成对应GT的box信息\n        bboxes_out = self.dboxes.clone().to(device=device)\n        bboxes_out[masks, :] = bboxes_in[best_dbox_idx[masks], :]\n\n        # Transform format to xywh format\n        x = 0.5 * (bboxes_out[:, 0] + bboxes_out[:, 2])  # x\n        y = 0.5 * (bboxes_out[:, 1] + bboxes_out[:, 3])  # y\n        w = bboxes_out[:, 2] - bboxes_out[:, 0]  # w\n        h = bboxes_out[:, 3] - bboxes_out[:, 1]  # h\n        bboxes_out[:, 0] = x\n        bboxes_out[:, 1] = y\n        bboxes_out[:, 2] = w\n        bboxes_out[:, 3] = h\n        return bboxes_out, labels_out\n\n    def scale_back_batch(self, bboxes_in, scores_in):\n        \"\"\"\n            将box格式从xywh转换回ltrb, 将预测目标score通过softmax处理\n            Do scale and transform from xywh to ltrb\n            suppose input N x 4 x num_bbox | N x label_num x num_bbox\n\n            bboxes_in: 是网络预测的xywh回归参数\n            scores_in: 是预测的每个default box的各目标概率\n        \"\"\"\n        if bboxes_in.device == torch.device(\"cpu\"):\n            self.dboxes = self.dboxes.cpu()\n            self.dboxes_xywh = self.dboxes_xywh.cpu()\n        else:\n            self.dboxes = self.dboxes.cuda()\n            self.dboxes_xywh = self.dboxes_xywh.cuda()\n\n        # Returns a view of the original tensor with its dimensions permuted.\n        bboxes_in = bboxes_in.permute(0, 2, 1)\n        scores_in = scores_in.permute(0, 2, 1)\n        # print(bboxes_in.is_contiguous())\n\n        bboxes_in[:, :, :2] = self.scale_xy * bboxes_in[:, :, :2]   # 预测的x, y回归参数\n        bboxes_in[:, :, 2:] = self.scale_wh * bboxes_in[:, :, 2:]   # 预测的w, h回归参数\n\n        # 将预测的回归参数叠加到default box上得到最终的预测边界框\n        bboxes_in[:, :, :2] = bboxes_in[:, :, :2] * self.dboxes_xywh[:, :, 2:] + self.dboxes_xywh[:, :, :2]\n        bboxes_in[:, :, 2:] = bboxes_in[:, :, 2:].exp() * self.dboxes_xywh[:, :, 2:]\n\n        # transform format to ltrb\n        l = bboxes_in[:, :, 0] - 0.5 * bboxes_in[:, :, 2]\n        t = bboxes_in[:, :, 1] - 0.5 * bboxes_in[:, :, 3]\n        r = bboxes_in[:, :, 0] + 0.5 * bboxes_in[:, :, 2]\n        b = bboxes_in[:, :, 1] + 0.5 * bboxes_in[:, :, 3]\n\n        bboxes_in[:, :, 0] = l  # xmin\n        bboxes_in[:, :, 1] = t  # ymin\n        bboxes_in[:, :, 2] = r  # xmax\n        bboxes_in[:, :, 3] = b  # ymax\n\n        return bboxes_in, F.softmax(scores_in, dim=-1)\n\n    def decode_batch(self, bboxes_in, scores_in, criteria=0.45, max_output=200):\n        # 将box格式从xywh转换回ltrb（方便后面非极大值抑制时求iou）, 将预测目标score通过softmax处理\n        bboxes, probs = self.scale_back_batch(bboxes_in, scores_in)\n\n        outputs = []\n        # 遍历一个batch中的每张image数据\n        for bbox, prob in zip(bboxes.split(1, 0), probs.split(1, 0)):\n            bbox = bbox.squeeze(0)\n            prob = prob.squeeze(0)\n            outputs.append(self.decode_single_new(bbox, prob, criteria, max_output))\n        return outputs\n\n    def decode_single_new(self, bboxes_in, scores_in, criteria, num_output=200):\n        \"\"\"\n        decode:\n            input  : bboxes_in (Tensor 8732 x 4), scores_in (Tensor 8732 x nitems)\n            output : bboxes_out (Tensor nboxes x 4), labels_out (Tensor nboxes)\n            criteria : IoU threshold of bboexes\n            max_output : maximum number of output bboxes\n        \"\"\"\n        device = bboxes_in.device\n        num_classes = scores_in.shape[-1]\n\n        # 对越界的bbox进行裁剪\n        bboxes_in = bboxes_in.clamp(min=0, max=1)\n\n        # [8732, 4] -> [8732, 21, 4]\n        bboxes_in = bboxes_in.repeat(1, num_classes).reshape(scores_in.shape[0], -1, 4)\n\n        # create labels for each prediction\n        labels = torch.arange(num_classes, device=device)\n        labels = labels.view(1, -1).expand_as(scores_in)\n\n        # remove prediction with the background label\n        # 移除归为背景类别的概率信息\n        bboxes_in = bboxes_in[:, 1:, :]\n        scores_in = scores_in[:, 1:]\n        labels = labels[:, 1:]\n\n        # batch everything, by making every class prediction be a separate instance\n        bboxes_in = bboxes_in.reshape(-1, 4)\n        scores_in = scores_in.reshape(-1)\n        labels = labels.reshape(-1)\n\n        # remove low scoring boxes\n        # 移除低概率目标，self.scores_thresh=0.05\n        inds = torch.nonzero(scores_in > 0.05, as_tuple=False).squeeze(1).to(device=device)\n        bboxes_in, scores_in, labels = bboxes_in[inds], scores_in[inds], labels[inds]\n\n        # remove empty boxes\n        ws, hs = bboxes_in[:, 2] - bboxes_in[:, 0], bboxes_in[:, 3] - bboxes_in[:, 1]\n        keep = (ws >= 0.1 / 300) & (hs >= 0.1 / 300)\n        keep = keep.nonzero(as_tuple=False).squeeze(1)\n        bboxes_in, scores_in, labels = bboxes_in[keep], scores_in[keep], labels[keep]\n\n        # non-maximum suppression\n        keep = batched_nms(bboxes_in, scores_in, labels, iou_threshold=criteria)\n\n        # keep only topk scoring predictions\n        keep = keep[:num_output]\n        bboxes_out = bboxes_in[keep, :]\n        scores_out = scores_in[keep]\n        labels_out = labels[keep]\n\n        return bboxes_out, labels_out, scores_out\n\n    # perform non-maximum suppression\n    def decode_single(self, bboxes_in, scores_in, criteria, max_output, max_num=200):\n        \"\"\"\n        decode:\n            input  : bboxes_in (Tensor 8732 x 4), scores_in (Tensor 8732 x nitems)\n            output : bboxes_out (Tensor nboxes x 4), labels_out (Tensor nboxes)\n            criteria : IoU threshold of bboexes\n            max_output : maximum number of output bboxes\n        \"\"\"\n        # Reference to https://github.com/amdegroot/ssd.pytorch\n        bboxes_out = []\n        scores_out = []\n        labels_out = []\n\n        # 非极大值抑制算法\n        # scores_in (Tensor 8732 x nitems), 遍历返回每一列数据，即8732个目标的同一类别的概率\n        for i, score in enumerate(scores_in.split(1, 1)):\n            # skip background\n            if i == 0:\n                continue\n\n            # [8732, 1] -> [8732]\n            score = score.squeeze(1)\n\n            # 虑除预测概率小于0.05的目标\n            mask = score > 0.05\n            bboxes, score = bboxes_in[mask, :], score[mask]\n            if score.size(0) == 0:\n                continue\n\n            # 按照分数从小到大排序\n            score_sorted, score_idx_sorted = score.sort(dim=0)\n\n            # select max_output indices\n            score_idx_sorted = score_idx_sorted[-max_num:]\n            candidates = []\n\n            while score_idx_sorted.numel() > 0:\n                idx = score_idx_sorted[-1].item()\n                # 获取排名前score_idx_sorted名的bboxes信息 Tensor:[score_idx_sorted, 4]\n                bboxes_sorted = bboxes[score_idx_sorted, :]\n                # 获取排名第一的bboxes信息 Tensor:[4]\n                bboxes_idx = bboxes[idx, :].unsqueeze(dim=0)\n                # 计算前score_idx_sorted名的bboxes与第一名的bboxes的iou\n                iou_sorted = calc_iou_tensor(bboxes_sorted, bboxes_idx).squeeze()\n\n                # we only need iou < criteria\n                # 丢弃与第一名iou > criteria的所有目标(包括自己本身)\n                score_idx_sorted = score_idx_sorted[iou_sorted < criteria]\n                # 保存第一名的索引信息\n                candidates.append(idx)\n\n            # 保存该类别通过非极大值抑制后的目标信息\n            bboxes_out.append(bboxes[candidates, :])   # bbox坐标信息\n            scores_out.append(score[candidates])       # score信息\n            labels_out.extend([i] * len(candidates))   # 标签信息\n\n        if not bboxes_out:  # 如果为空的话，返回空tensor，注意boxes对应的空tensor size，防止验证时出错\n            return [torch.empty(size=(0, 4)), torch.empty(size=(0,), dtype=torch.int64), torch.empty(size=(0,))]\n\n        bboxes_out = torch.cat(bboxes_out, dim=0).contiguous()\n        scores_out = torch.cat(scores_out, dim=0).contiguous()\n        labels_out = torch.as_tensor(labels_out, dtype=torch.long).to(device=device)\n\n        # 对所有目标的概率进行排序（无论是什 么类别）,取前max_num个目标\n        _, max_ids = scores_out.sort(dim=0)\n        max_ids = max_ids[-max_output:]\n        return bboxes_out[max_ids, :], labels_out[max_ids], scores_out[max_ids]\n\n\n# figsize = 300  # 输入网络的图像大小\n# feat_size = [38, 19, 10, 5, 3, 1]   # 每个预测层的feature map尺寸\n# steps = [8, 16, 32, 64, 100, 300]   # 每个特征层上的一个cell在原图上的跨度\n# scales = [21, 45, 99, 153, 207, 261, 315]  # 每个特征层上预测的default box的scale\n# aspect_ratios = [[2], [2, 3], [2, 3], [2, 3], [2], [2]]  # 每个预测特征层上预测的default box的ratios\nclass DefaultBoxes(object):\n    def __init__(self, fig_size, feat_size, steps, scales, aspect_ratios, scale_xy=0.1, scale_wh=0.2):\n        self.fig_size = fig_size   # 输入网络的图像大小 300\n        # [38, 19, 10, 5, 3, 1]\n        self.feat_size = feat_size  # 每个预测层的feature map尺寸\n\n        self.scale_xy_ = scale_xy\n\n        self.scale_wh_ = scale_wh\n\n        # [8, 16, 32, 64, 100, 300]\n        self.steps = steps    # 每个特征层上的一个cell在原图上的跨度\n\n        # [21, 45, 99, 153, 207, 261, 315]\n        self.scales = scales  # 每个特征层上预测的default box的scale\n\n        fk = fig_size / np.array(steps)     # 计算每层特征层的fk\n\n        # [[2], [2, 3], [2, 3], [2, 3], [2], [2]]\n        self.aspect_ratios = aspect_ratios  # 每个预测特征层上预测的default box的ratios\n\n        self.default_boxes = []\n        # size of feature and number of feature\n        # 遍历每层特征层，计算default box\n        for idx, sfeat in enumerate(self.feat_size):\n            sk1 = scales[idx] / fig_size  # scale转为相对值[0-1]\n            sk2 = scales[idx + 1] / fig_size  # scale转为相对值[0-1]\n            sk3 = sqrt(sk1 * sk2)\n            # 先添加两个1:1比例的default box宽和高\n            all_sizes = [(sk1, sk1), (sk3, sk3)]\n\n            # 再将剩下不同比例的default box宽和高添加到all_sizes中\n            for alpha in aspect_ratios[idx]:\n                w, h = sk1 * sqrt(alpha), sk1 / sqrt(alpha)\n                all_sizes.append((w, h))\n                all_sizes.append((h, w))\n\n            # 计算当前特征层对应原图上的所有default box\n            for w, h in all_sizes:\n                for i, j in itertools.product(range(sfeat), repeat=2):  # i -> 行（y）， j -> 列（x）\n                    # 计算每个default box的中心坐标（范围是在0-1之间）\n                    cx, cy = (j + 0.5) / fk[idx], (i + 0.5) / fk[idx]\n                    self.default_boxes.append((cx, cy, w, h))\n\n        # 将default_boxes转为tensor格式\n        self.dboxes = torch.as_tensor(self.default_boxes, dtype=torch.float32)  # 这里不转类型会报错\n        self.dboxes.clamp_(min=0, max=1)  # 将坐标（x, y, w, h）都限制在0-1之间\n\n        # For IoU calculation\n        # ltrb is left top coordinate and right bottom coordinate\n        # 将(x, y, w, h)转换成(xmin, ymin, xmax, ymax)，方便后续计算IoU(匹配正负样本时)\n        self.dboxes_ltrb = self.dboxes.clone()\n        self.dboxes_ltrb[:, 0] = self.dboxes[:, 0] - 0.5 * self.dboxes[:, 2]   # xmin\n        self.dboxes_ltrb[:, 1] = self.dboxes[:, 1] - 0.5 * self.dboxes[:, 3]   # ymin\n        self.dboxes_ltrb[:, 2] = self.dboxes[:, 0] + 0.5 * self.dboxes[:, 2]   # xmax\n        self.dboxes_ltrb[:, 3] = self.dboxes[:, 1] + 0.5 * self.dboxes[:, 3]   # ymax\n\n    @property\n    def scale_xy(self):\n        return self.scale_xy_\n\n    @property\n    def scale_wh(self):\n        return self.scale_wh_\n\n    def __call__(self, order='ltrb'):\n        # 根据需求返回对应格式的default box\n        if order == 'ltrb':\n            return self.dboxes_ltrb\n\n        if order == 'xywh':\n            return self.dboxes\n\n\ndef dboxes300_coco():\n    figsize = 300  # 输入网络的图像大小\n    feat_size = [38, 19, 10, 5, 3, 1]   # 每个预测层的feature map尺寸\n    steps = [8, 16, 32, 64, 100, 300]   # 每个特征层上的一个cell在原图上的跨度\n    scales = [21, 45, 99, 153, 207, 261, 315]  # 每个特征层上预测的default box的scale\n    aspect_ratios = [[2], [2, 3], [2, 3], [2, 3], [2], [2]]  # 每个预测特征层上预测的default box的ratios\n    dboxes = DefaultBoxes(figsize, feat_size, steps, scales, aspect_ratios)\n    return dboxes\n\n\ndef nms(boxes, scores, iou_threshold):\n    # type: (Tensor, Tensor, float) -> Tensor\n    \"\"\"\n    Performs non-maximum suppression (NMS) on the boxes according\n    to their intersection-over-union (IoU).\n\n    NMS iteratively removes lower scoring boxes which have an\n    IoU greater than iou_threshold with another (higher scoring)\n    box.\n\n    Parameters\n    ----------\n    boxes : Tensor[N, 4])\n        boxes to perform NMS on. They\n        are expected to be in (x1, y1, x2, y2) format\n    scores : Tensor[N]\n        scores for each one of the boxes\n    iou_threshold : float\n        discards all overlapping\n        boxes with IoU < iou_threshold\n\n    Returns\n    -------\n    keep : Tensor\n        int64 tensor with the indices\n        of the elements that have been kept\n        by NMS, sorted in decreasing order of scores\n    \"\"\"\n    return torch.ops.torchvision.nms(boxes, scores, iou_threshold)\n\n\ndef batched_nms(boxes, scores, idxs, iou_threshold):\n    # type: (Tensor, Tensor, Tensor, float) -> Tensor\n    \"\"\"\n    Performs non-maximum suppression in a batched fashion.\n\n    Each index value correspond to a category, and NMS\n    will not be applied between elements of different categories.\n\n    Parameters\n    ----------\n    boxes : Tensor[N, 4]\n        boxes where NMS will be performed. They\n        are expected to be in (x1, y1, x2, y2) format\n    scores : Tensor[N]\n        scores for each one of the boxes\n    idxs : Tensor[N]\n        indices of the categories for each one of the boxes.\n    iou_threshold : float\n        discards all overlapping boxes\n        with IoU < iou_threshold\n\n    Returns\n    -------\n    keep : Tensor\n        int64 tensor with the indices of\n        the elements that have been kept by NMS, sorted\n        in decreasing order of scores\n    \"\"\"\n    if boxes.numel() == 0:\n        return torch.empty((0,), dtype=torch.int64, device=boxes.device)\n\n    # strategy: in order to perform NMS independently per class.\n    # we add an offset to all the boxes. The offset is dependent\n    # only on the class idx, and is large enough so that boxes\n    # from different classes do not overlap\n    # 获取所有boxes中最大的坐标值（xmin, ymin, xmax, ymax）\n    max_coordinate = boxes.max()\n\n    # to(): Performs Tensor dtype and/or device conversion\n    # 为每一个类别生成一个很大的偏移量\n    # 这里的to只是让生成tensor的dytpe和device与boxes保持一致\n    offsets = idxs.to(boxes) * (max_coordinate + 1)\n    # boxes加上对应层的偏移量后，保证不同类别之间boxes不会有重合的现象\n    boxes_for_nms = boxes + offsets[:, None]\n    keep = nms(boxes_for_nms, scores, iou_threshold)\n    return keep\n\n\nclass PostProcess(nn.Module):\n    def __init__(self, dboxes):\n        super(PostProcess, self).__init__()\n        # [num_anchors, 4] -> [1, num_anchors, 4]\n        self.dboxes_xywh = nn.Parameter(dboxes(order='xywh').unsqueeze(dim=0),\n                                        requires_grad=False)\n        self.scale_xy = dboxes.scale_xy  # 0.1\n        self.scale_wh = dboxes.scale_wh  # 0.2\n\n        self.criteria = 0.5\n        self.max_output = 100\n\n    def scale_back_batch(self, bboxes_in, scores_in):\n        # type: (Tensor, Tensor) -> Tuple[Tensor, Tensor]\n        \"\"\"\n            1）通过预测的boxes回归参数得到最终预测坐标\n            2）将box格式从xywh转换回ltrb\n            3）将预测目标score通过softmax处理\n            Do scale and transform from xywh to ltrb\n            suppose input N x 4 x num_bbox | N x label_num x num_bbox\n\n            bboxes_in: [N, 4, 8732]是网络预测的xywh回归参数\n            scores_in: [N, label_num, 8732]是预测的每个default box的各目标概率\n        \"\"\"\n\n        # Returns a view of the original tensor with its dimensions permuted.\n        # [batch, 4, 8732] -> [batch, 8732, 4]\n        bboxes_in = bboxes_in.permute(0, 2, 1)\n        # [batch, label_num, 8732] -> [batch, 8732, label_num]\n        scores_in = scores_in.permute(0, 2, 1)\n        # print(bboxes_in.is_contiguous())\n\n        bboxes_in[:, :, :2] = self.scale_xy * bboxes_in[:, :, :2]   # 预测的x, y回归参数\n        bboxes_in[:, :, 2:] = self.scale_wh * bboxes_in[:, :, 2:]   # 预测的w, h回归参数\n\n        # 将预测的回归参数叠加到default box上得到最终的预测边界框\n        bboxes_in[:, :, :2] = bboxes_in[:, :, :2] * self.dboxes_xywh[:, :, 2:] + self.dboxes_xywh[:, :, :2]\n        bboxes_in[:, :, 2:] = bboxes_in[:, :, 2:].exp() * self.dboxes_xywh[:, :, 2:]\n\n        # transform format to ltrb\n        l = bboxes_in[:, :, 0] - 0.5 * bboxes_in[:, :, 2]\n        t = bboxes_in[:, :, 1] - 0.5 * bboxes_in[:, :, 3]\n        r = bboxes_in[:, :, 0] + 0.5 * bboxes_in[:, :, 2]\n        b = bboxes_in[:, :, 1] + 0.5 * bboxes_in[:, :, 3]\n\n        bboxes_in[:, :, 0] = l  # xmin\n        bboxes_in[:, :, 1] = t  # ymin\n        bboxes_in[:, :, 2] = r  # xmax\n        bboxes_in[:, :, 3] = b  # ymax\n\n        # scores_in: [batch, 8732, label_num]\n        return bboxes_in, F.softmax(scores_in, dim=-1)\n\n    def decode_single_new(self, bboxes_in, scores_in, criteria, num_output):\n        # type: (Tensor, Tensor, float, int) -> Tuple[Tensor, Tensor, Tensor]\n        \"\"\"\n        decode:\n            input  : bboxes_in (Tensor 8732 x 4), scores_in (Tensor 8732 x nitems)\n            output : bboxes_out (Tensor nboxes x 4), labels_out (Tensor nboxes)\n            criteria : IoU threshold of bboexes\n            max_output : maximum number of output bboxes\n        \"\"\"\n        device = bboxes_in.device\n        num_classes = scores_in.shape[-1]\n\n        # 对越界的bbox进行裁剪\n        bboxes_in = bboxes_in.clamp(min=0, max=1)\n\n        # [8732, 4] -> [8732, 21, 4]\n        bboxes_in = bboxes_in.repeat(1, num_classes).reshape(scores_in.shape[0], -1, 4)\n\n        # create labels for each prediction\n        labels = torch.arange(num_classes, device=device)\n        # [num_classes] -> [8732, num_classes]\n        labels = labels.view(1, -1).expand_as(scores_in)\n\n        # remove prediction with the background label\n        # 移除归为背景类别的概率信息\n        bboxes_in = bboxes_in[:, 1:, :]  # [8732, 21, 4] -> [8732, 20, 4]\n        scores_in = scores_in[:, 1:]  # [8732, 21] -> [8732, 20]\n        labels = labels[:, 1:]  # [8732, 21] -> [8732, 20]\n\n        # batch everything, by making every class prediction be a separate instance\n        bboxes_in = bboxes_in.reshape(-1, 4)  # [8732, 20, 4] -> [8732x20, 4]\n        scores_in = scores_in.reshape(-1)  # [8732, 20] -> [8732x20]\n        labels = labels.reshape(-1)  # [8732, 20] -> [8732x20]\n\n        # remove low scoring boxes\n        # 移除低概率目标，self.scores_thresh=0.05\n        # inds = torch.nonzero(scores_in > 0.05).squeeze(1)\n        inds = torch.where(torch.gt(scores_in, 0.05))[0]\n        bboxes_in, scores_in, labels = bboxes_in[inds, :], scores_in[inds], labels[inds]\n\n        # remove empty boxes\n        ws, hs = bboxes_in[:, 2] - bboxes_in[:, 0], bboxes_in[:, 3] - bboxes_in[:, 1]\n        keep = (ws >= 1 / 300) & (hs >= 1 / 300)\n        # keep = keep.nonzero().squeeze(1)\n        keep = torch.where(keep)[0]\n        bboxes_in, scores_in, labels = bboxes_in[keep], scores_in[keep], labels[keep]\n\n        # non-maximum suppression\n        keep = batched_nms(bboxes_in, scores_in, labels, iou_threshold=criteria)\n\n        # keep only topk scoring predictions\n        keep = keep[:num_output]\n        bboxes_out = bboxes_in[keep, :]\n        scores_out = scores_in[keep]\n        labels_out = labels[keep]\n\n        return bboxes_out, labels_out, scores_out\n\n    def forward(self, bboxes_in, scores_in):\n        # 通过预测的boxes回归参数得到最终预测坐标, 将预测目标score通过softmax处理\n        bboxes, probs = self.scale_back_batch(bboxes_in, scores_in)\n\n        outputs = torch.jit.annotate(List[Tuple[Tensor, Tensor, Tensor]], [])\n        # 遍历一个batch中的每张image数据\n        # bboxes: [batch, 8732, 4]\n        for bbox, prob in zip(bboxes.split(1, 0), probs.split(1, 0)):  # split_size, split_dim\n            # bbox: [1, 8732, 4]\n            bbox = bbox.squeeze(0)\n            prob = prob.squeeze(0)\n            outputs.append(self.decode_single_new(bbox, prob, self.criteria, self.max_output))\n        return outputs","metadata":{"_kg_hide-input":false,"execution":{"iopub.status.busy":"2023-08-10T05:46:26.786238Z","iopub.execute_input":"2023-08-10T05:46:26.787181Z","iopub.status.idle":"2023-08-10T05:46:26.902084Z","shell.execute_reply.started":"2023-08-10T05:46:26.787143Z","shell.execute_reply":"2023-08-10T05:46:26.900565Z"},"trusted":true},"execution_count":7,"outputs":[]},{"cell_type":"code","source":"# src/ssd_model.py\n\nimport torch\nfrom torch import nn, Tensor\nfrom torch.jit.annotations import List\n\nfrom torchvision import models\n\n# from .res50_backbone import resnet50\n# from .utils import dboxes300_coco, Encoder, PostProcess\n\n\nclass Backbone(nn.Module):\n    \"\"\"\n        定义一个backbone\n    \"\"\"\n\n    def __init__(self, pretrain_path=None):\n        super(Backbone, self).__init__()\n\n        # 定义一个resnet50\n        net = models.resnet50(pretrained=True, progress=True)\n\n        # 后续的通道数\n        self.out_channels = [2048, 1024, 1024, 512, 512, 512]\n\n        # 加载预训练模型\n        if pretrain_path is not None:\n            net.load_state_dict(torch.load(pretrain_path))\n\n        # 截取特征提出部分\n        self.feature_extractor = nn.Sequential(*list(net.children())[:7])\n\n        # 修改其中的属性\n        conv4_block1 = self.feature_extractor[-1][0]\n        # 修改conv4_block1的stride，从2->1\n        conv4_block1.conv1.stride = (1, 1)\n        conv4_block1.conv2.stride = (1, 1)\n        conv4_block1.downsample[0].stride = (1, 1)\n\n    def forward(self, x):\n        x = self.feature_extractor(x)\n        return x\n    \n# 主模型\nclass SSD300(nn.Module):\n    \"\"\"\n        SSD300主模型\n    \"\"\"\n\n    def __init__(self, backbone=None, num_classes=21):\n        \"\"\"\n            设置分类的数量\n\n        \"\"\"\n        super(SSD300, self).__init__()\n\n        # 参数校验\n        if backbone is None:\n            raise Exception(\"backbone is None\")\n\n        if not hasattr(backbone, \"out_channels\"):\n            raise Exception(\"the backbone not has attribute: out_channel\")\n\n        self.feature_extractor = backbone\n\n        self.num_classes = num_classes\n        \n        # [b, 1024, 76, 76] -> [b, 2048, 38, 38]\n        self.conv = nn.Conv2d(in_channels=1024, out_channels=2048, kernel_size=3, padding=1, stride=2)\n\n        # 构建自定义的特征层\n        # out_channels = [2048, 1024, 1024, 512, 512, 512] for resnet50\n        self._build_additional_features(self.feature_extractor.out_channels)\n        \n        # 每个特征层上每个特征点对应的锚框数量\n        self.num_defaults = [4, 6, 6, 6, 4, 4]\n\n        self.belt_pos = nn.Conv2d(in_channels=512, out_channels=2, kernel_size=1, padding=0)\n        \n        location_extractors = []\n        confidence_extractors = []\n\n        for nd, oc in zip(self.num_defaults, self.feature_extractor.out_channels):\n            # nd is number_default_boxes, oc is output_channel\n            location_extractors.append(nn.Conv2d(oc, nd * 4, kernel_size=3, padding=1))\n            confidence_extractors.append(nn.Conv2d(oc, nd * self.num_classes, kernel_size=3, padding=1))\n\n        self.loc = nn.ModuleList(location_extractors)\n        self.conf = nn.ModuleList(confidence_extractors)\n        self._init_weights()\n\n        # 模型锚框的生成策略\n        default_box = dboxes300_coco()\n\n        # 下面三个是核心\n        self.compute_loss = Loss(default_box)\n        self.encoder = Encoder(default_box)\n        self.postprocess = PostProcess(default_box)\n\n    def _build_additional_features(self, input_size):\n        \"\"\"\n        为backbone(resnet50)添加额外的一系列卷积层，得到相应的一系列特征提取器\n        :param input_size:\n        :return:\n        \"\"\"\n        additional_blocks = []\n        # input_size = [2048, 1024, 1024, 512, 512, 512] for resnet50\n        # shape :          76   [38,   19,   10,   5,   3]\n        # input_size[:-1]：1024 [2048, 1024, 1024, 512, 512]\n        # middle_channels：     [512,  512,  256,  256, 256]\n        # input_size[1:]：      [1024, 1024, 512,  512, 512]\n        middle_channels = [512, 512, 256, 256, 256]\n        for i, (input_ch, output_ch, middle_ch) in enumerate(zip(input_size[:-1], input_size[1:], middle_channels)):\n            padding, stride = (1, 2) if i < 3 else (0, 1)\n            layer = nn.Sequential(\n                nn.Conv2d(input_ch, middle_ch, kernel_size=1, bias=False),\n                nn.BatchNorm2d(middle_ch),\n                nn.ReLU(inplace=True),\n                nn.Conv2d(middle_ch, output_ch, kernel_size=3, padding=padding, stride=stride, bias=False),\n                nn.BatchNorm2d(output_ch),\n#                 nn.ReLU(inplace=True),\n            )\n            additional_blocks.append(layer)\n            \n        self.additional_blocks = nn.ModuleList(additional_blocks)\n        \n        shortcut_blocks = [\n            # [b, 1024, 76, 76] -> [b, 1024, 19, 19]\n            nn.Conv2d(in_channels=1024, out_channels=1024, kernel_size=1, stride=4),\n            # [b, 2048, 38, 38] -> [b, 1024, 10, 10]\n            nn.Conv2d(in_channels=2048, out_channels=1024, kernel_size=1, stride=4),\n            # [b, 1024, 19, 19] -> [b, 512, 5, 5]\n            nn.Conv2d(in_channels=1024, out_channels=512, kernel_size=1, stride=4),\n            # [b, 1024, 10, 10] -> [b, 512, 3, 3]\n            nn.Conv2d(in_channels=1024, out_channels=512, kernel_size=1, stride=4),\n            # [b, 512, 5, 5] -> [b, 512, 3, 3]\n            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=1, stride=2)\n        ]\n        \n        self.shortcut_blocks = nn.ModuleList(shortcut_blocks)\n\n    def _init_weights(self):\n        \"\"\"\n            要不要自己去初始化权重？？？\n        \"\"\"\n        layers = [self.conv, *self.additional_blocks, *self.shortcut_blocks, *self.loc, *self.conf]\n        for layer in layers:\n            for param in layer.parameters():\n                if param.dim() > 1:\n                    nn.init.xavier_uniform_(param)\n\n    # Shape the classifier to the view of bboxes\n    def bbox_view(self, features, loc_extractor, conf_extractor):\n        locs = []\n        confs = []\n        for f, l, c in zip(features, loc_extractor, conf_extractor):\n            # [batch, n*4, feat_size, feat_size] -> [batch, 4, -1]\n            locs.append(l(f).view(f.size(0), 4, -1))\n            # [batch, n*classes, feat_size, feat_size] -> [batch, classes, -1]\n            confs.append(c(f).view(f.size(0), self.num_classes, -1))\n\n        locs, confs = torch.cat(locs, 2).contiguous(), torch.cat(confs, 2).contiguous()\n        return locs, confs\n\n    def forward(self, eopch, image, targets=None):\n        # [b, 3, 600, 600] -> [b, 1024, 76, 76]\n        x = self.feature_extractor(image)\n        \n        shortcuts = []\n        shortcuts.append(self.shortcut_blocks[0](x))\n        \n        # [b, 1024, 76, 76] -> [b, 2048, 38, 38]\n        x = self.conv(x)\n\n        # Feature Map 2048, 19x19x1024, 10x10x1024, 5x5x512, 3x3x512, 1x1x512\n        detection_features = torch.jit.annotate(List[Tensor], [])  # [x]\n        detection_features.append(x)\n            \n        relu_inplace = nn.ReLU(inplace=True)\n        relu = nn.ReLU(inplace=False)\n        \n        additional_blocks_count = len(self.additional_blocks)\n        for i in range(additional_blocks_count):\n            # for sub_layer in layer.children():\n            #    try:\n            #        print(f\"{sub_layer}, weight.grad:{sub_layer.weight.grad}\")\n            #    except:\n            #        pass\n            \n            # shape          :  76   [38,   19,   10,   5,   3]\n            # input_size[:-1]： 1024 [2048, 1024, 1024, 512, 512]\n            # middle_channels：      [512,  512,  256,  256, 256]\n            # input_size[1:]：       [1024, 1024, 512,  512, 512]\n            if i < 4:\n                shortcuts.append(self.shortcut_blocks[i + 1](x))\n            \n            if i >= 1:\n                x = self.additional_blocks[i](relu_inplace(x + shortcuts[i - 1]))\n                if i == additional_blocks_count - 1:\n                    x = relu_inplace(x)\n            else:\n                x = relu_inplace(self.additional_blocks[i](x))\n        \n            detection_features.append(x)\n            \n        # [b, 512, 1, 1] -> [b, 2, 1, 1]\n        ptask2_label = self.belt_pos(x)\n        # [b, 2, 1, 1] -> [b, 2]\n        ptask2_label = ptask2_label.reshape(-1, 2)\n\n        # Feature Map 38x38x4, 19x19x6, 10x10x6, 5x5x6, 3x3x4, 1x1x4\n        locs, confs = self.bbox_view(detection_features, self.loc, self.conf)\n\n        # For SSD 300, shall return nbatch x 8732 x {nlabels, nlocs} results\n        # 38x38x4 + 19x19x6 + 10x10x6 + 5x5x6 + 3x3x4 + 1x1x4 = 8732\n\n        if self.training:\n            if targets is None:\n                raise ValueError(\"In training mode, targets should be passed\")\n            # bboxes_out (Tensor 8732 x 4), labels_out (Tensor 8732)\n            bboxes_out = targets['boxes']\n            bboxes_out = bboxes_out.transpose(1, 2).contiguous()\n            # print(bboxes_out.is_contiguous())\n            labels_out = targets['labels']\n            # print(labels_out.is_contiguous())\n\n#             if epoch % 20 == 0:\n#                 self.show_train_effect(image, locs, confs, bboxes_out, labels_out)\n            \n            # ploc, plabel, gloc, glabel\n            loss = self.compute_loss(locs, \n                                     confs, \n                                     bboxes_out, \n                                     labels_out, \n                                     ptask2_label, \n                                     targets['task2_label'])\n            return {\"total_losses\": loss}\n\n        # 将预测回归参数叠加到default box上得到最终预测box，并执行非极大值抑制虑除重叠框\n        # results = self.encoder.decode_batch(locs, confs)\n        results = {}\n        results[PREDICT_RESULT_TASK1] = self.postprocess(locs, confs)\n        results[PREDICT_RESULT_TASK2] = ptask2_label\n        return results\n    \n    def show_train_effect(self, imgs, locs_pred, confs_pred, bboxes_out, labels_out):\n        \n        batch_size = locs_pred.shape[0]\n        # 从当前批次任选一图片，展示预测的偏移准确度效果\n        choice = random.randint(0, batch_size - 1)\n        # [32, 4, 8732], [32, 3, 8732] -> [4, 8732], [3, 8732]\n        locs_pred, confs_pred = locs_pred[choice], confs_pred[choice]\n        # [32, 4, 8732], [32, 8732] -> [4, 8732], [8732] \n        bboxes_out, labels_out = bboxes_out[choice], labels_out[choice]\n        \n        # [4, 8732] -> [8732, 4]\n        locs_pred = locs_pred.permute(1, 0)\n        # [3, 8732] -> [8732, 3]\n        confs_pred = confs_pred.permute(1, 0)\n        # [4, 8732] -> [8732, 4]\n        bboxes_out = bboxes_out.permute(1, 0)\n        \n        # [8732] -> [8732]\n        true_flag = labels_out > 0\n        \n        # [8732, 4] -> [N', 4]\n        locs_pred = locs_pred[true_flag]\n        # [8732, 4] -> [N']\n        confs_pred = confs_pred[true_flag]\n        \n        pred_scores, confs_pred = torch.max(input=confs_pred, dim=1)\n        # [8732, 4] -> [N', 4]\n        bboxes_out = bboxes_out[true_flag]\n        # [8732] -> [N']\n        labels_out = labels_out[true_flag]\n        \n        bboxes = torch.zeros_like(bboxes_out).to(device=device)\n        bboxes[:, 0] = bboxes_out[:, 0] - bboxes_out[:, 2] * 0.5\n        bboxes[:, 1] = bboxes_out[:, 1] - bboxes_out[:, 3] * 0.5\n        bboxes[:, 2] = bboxes_out[:, 0] + bboxes_out[:, 2] * 0.5\n        bboxes[:, 3] = bboxes_out[:, 1] + bboxes_out[:, 3] * 0.5\n        \n        dboxes_xywh = dboxes300_coco()(order='xywh').to(device=device)\n        dboxes_xywh = dboxes_xywh[true_flag]\n        locs_pred[:, :2] = 0.1 * locs_pred[:, :2]   # 预测的x, y回归参数\n        locs_pred[:, 2:] = 0.2 * locs_pred[:, 2:]   # 预测的w, h回归参数\n\n        # 将预测的回归参数叠加到default box上得到最终的预测边界框\n        locs_pred[:, :2] = locs_pred[:, :2] * dboxes_xywh[:, 2:] + dboxes_xywh[:, :2]\n        locs_pred[:, 2:] = locs_pred[:, 2:].exp() * dboxes_xywh[:, 2:]\n\n        # transform format to ltrb\n        l = locs_pred[:, 0] - 0.5 * locs_pred[:, 2]\n        t = locs_pred[:, 1] - 0.5 * locs_pred[:, 3]\n        r = locs_pred[:, 0] + 0.5 * locs_pred[:, 2]\n        b = locs_pred[:, 1] + 0.5 * locs_pred[:, 3]\n\n        locs_pred[:, 0] = l  # xmin\n        locs_pred[:, 1] = t  # ymin\n        locs_pred[:, 2] = r  # xmax\n        locs_pred[:, 3] = b  # ymax\n        \n#         print(f\"train: bboxes:{bboxes[:5].detach()}, labels_out:{labels_out[:5].detach()}\")\n#         print(f\"train: locs_pred:{locs_pred[:5].detach()}, confs_pred:{labels_out[:5].detach()}, pred_scores:{pred_scores}\")\n        \n        save_anchor_pic(image=imgs, \n                        path=\"./output/train_pred.jpg\", \n                        bboxes=locs_pred.detach(), \n                        bboxes_label = confs_pred.detach(), \n                        bboxes_scores=pred_scores,\n                        height=600, \n                        width=600, \n                        targets=bboxes, \n                        targets_label = labels_out.detach(),\n                        choice=choice)\n\nclass Loss(nn.Module):\n    \"\"\"\n        Implements the loss as the sum of the followings:\n        1. Confidence Loss: All labels, with hard negative mining\n        2. Localization Loss: Only on positive labels\n        Suppose input dboxes has the shape 8732x4\n    \"\"\"\n\n    def __init__(self, dboxes):\n        super(Loss, self).__init__()\n        # Two factor are from following links\n        # http://jany.st/post/2017-11-05-single-shot-detector-ssd-from-scratch-in-tensorflow.html\n        self.scale_xy = 1.0 / dboxes.scale_xy  # 10\n        self.scale_wh = 1.0 / dboxes.scale_wh  # 5\n\n        self.location_loss = nn.SmoothL1Loss(reduction='none')\n        # [num_anchors, 4] -> [4, num_anchors] -> [1, 4, num_anchors]\n        self.dboxes = nn.Parameter(dboxes(order=\"xywh\").transpose(0, 1).unsqueeze(dim=0),\n                                   requires_grad=False)\n\n        self.confidence_loss = nn.CrossEntropyLoss(reduction='none')\n        self.task2_loss = nn.CrossEntropyLoss()\n\n    def _location_vec(self, loc):\n        # type: (Tensor) -> Tensor\n        \"\"\"\n        Generate Location Vectors\n        计算ground truth相对anchors的回归参数\n        :param loc: anchor匹配到的对应GTBOX Nx4x8732\n        :return:\n        \"\"\"\n        \n        gxy = self.scale_xy * (loc[:, :2, :] - self.dboxes[:, :2, :]) / self.dboxes[:, 2:, :]  # Nx2x8732\n        gwh = self.scale_wh * (loc[:, 2:, :] / self.dboxes[:, 2:, :]).log()  # Nx2x8732\n        return torch.cat((gxy, gwh), dim=1).contiguous()\n\n    def forward(self, ploc, plabel, gloc, glabel, ptask2_label, gtask2_label):\n        # type: (Tensor, Tensor, Tensor, Tensor, Tensor) -> Tensor\n        \"\"\"\n            ploc, plabel: Nx4x8732, Nxlabel_numx8732\n                predicted location and labels\n\n            gloc, glabel: Nx4x8732, Nx8732\n                ground truth location and labels\n            ptask2_label, gtask2_label: Nx2, Nx2\n        \"\"\"\n        # 获取正样本的mask  Tensor: [N, 8732]\n        mask = torch.gt(glabel, 0)  # (gt: >)\n        # mask1 = torch.nonzero(glabel)\n        # 计算一个batch中的每张图片的正样本个数 Tensor: [N]\n        pos_num = mask.sum(dim=1)\n\n        # 计算gt的location回归参数 Tensor: [N, 4, 8732]\n        vec_gd = self._location_vec(gloc)\n\n        # sum on four coordinates, and mask\n        # 计算定位损失(只有正样本)\n        loc_loss = self.location_loss(ploc, vec_gd).sum(dim=1)  # Tensor: [N, 8732]\n        loc_loss = (mask.float() * loc_loss).sum(dim=1)  # Tenosr: [N]\n        \n        # hard negative mining Tenosr: [N, 8732]\n        con = self.confidence_loss(plabel, glabel)\n\n        # positive mask will never selected\n        # 获取负样本\n        con_neg = con.clone()\n        con_neg[mask] = 0.0\n        # 按照confidence_loss降序排列 con_idx(Tensor: [N, 8732])\n        _, con_idx = con_neg.sort(dim=1, descending=True)\n        # 第一次降序sort，置信度越高者，排在越前面,\n        # 第二次对con_idx做升序排列，则排序结果又恢复为con_neg中的原有排序顺序，而每个位置对应在con_rank的值，\n        # 记录的是其在con_idx的排序顺序，也就是元素序列值越小，confidence值越大，这样就可以直接mask con_neg，\n        # 获取指定位置的元素, 比如：\n        # >>> mask = torch.as_tensor([[True, False, False, True], [True, False, False, False]])\n        # >>> mask\n        # tensor([[True, False, False, True],\n        #         [True, False, False, False]])\n        # >>> con[mask]\n        # tensor([0.2000, 0.3400, 0.5000])\n        # >>> con_neg = con.clone()\n        # >>> con_neg\n        # tensor([[0.2000, 0.3000, 0.1200, 0.3400],\n        #         [0.5000, 0.2000, 0.6000, 0.7000]])\n        # >>> con_neg = con[mask]\n        # >>> con_neg\n        # tensor([0.2000, 0.3400, 0.5000])\n        # >>> con_neg = con.clone()\n        # >>> con_neg\n        # tensor([[0.2000, 0.3000, 0.1200, 0.3400],\n        #         [0.5000, 0.2000, 0.6000, 0.7000]])\n        # >>> con_neg[mask] = 0\n        # >>> con_neg\n        # tensor([[0.0000, 0.3000, 0.1200, 0.0000],\n        #         [0.0000, 0.2000, 0.6000, 0.7000]])\n        # >>> _, con_idx = torch.sort(con_neg, descending=True)\n        # >>> _\n        # tensor([[0.3000, 0.1200, 0.0000, 0.0000],\n        #         [0.7000, 0.6000, 0.2000, 0.0000]])\n        # >>> con_idx\n        # tensor([[1, 2, 0, 3],\n        #         [3, 2, 1, 0]])\n        # >>> _, con_rank = torch.sort(con_idx)\n        # >>> _\n        # tensor([[0, 1, 2, 3],\n        #         [0, 1, 2, 3]])\n        # >>> con_rank\n        # tensor([[2, 0, 1, 3],\n        #         [3, 2, 1, 0]])\n        # >>> con * torch.lt(con_rank, 2)\n        # tensor([[0.0000, 0.3000, 0.1200, 0.0000],\n        #         [0.0000, 0.0000, 0.6000, 0.7000]])\n        _, con_rank = con_idx.sort(dim=1)  # 这个步骤比较巧妙\n\n        # number of negative three times positive\n        # 用于损失计算的负样本数是正样本的3倍（在原论文Hard negative mining部分），\n        # 但不能超过总样本数8732\n        neg_num = torch.clamp(3 * pos_num, max=mask.size(1)).unsqueeze(-1)\n        neg_mask = torch.lt(con_rank, neg_num)  # (lt: <) Tensor [N, 8732]\n\n        # confidence最终loss使用选取的正样本loss+选取的负样本loss\n        con_loss = (con * (mask.float() + neg_mask.float())).sum(dim=1)  # Tensor [N]\n\n        # eg. [15, 3, 5, 0] -> [1.0, 1.0, 1.0, 0.0]\n        num_mask = torch.gt(pos_num, 0).float()  # 统计一个batch中的每张图像中是否存在正样本\n        pos_num = pos_num.float().clamp(min=1e-6)  # 防止出现分母为零的情况\n        \n        loc_loss = (loc_loss * num_mask / pos_num).mean(dim=0)  # 只计算存在正样本的图像损失\n        con_loss = (con_loss * num_mask / pos_num).mean(dim=0)  # 只计算存在正样本的图像损失\n        \n        task2_loss = self.task2_loss(ptask2_label, gtask2_label)\n        \n        # avoid no object detected\n        # 避免出现图像中没有GTBOX的情况\n        total_loss = 0.5 * (loc_loss + con_loss) + 0.5 * task2_loss\n        \n#         print(f\"---loss: con_loss:{con_loss:.4}, loc_loss:{loc_loss:.4}, task2_loss:{task2_loss:.4}, total_loss:{total_loss:.4}\")\n        \n        return total_loss","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:46:26.904697Z","iopub.execute_input":"2023-08-10T05:46:26.905592Z","iopub.status.idle":"2023-08-10T05:46:26.986317Z","shell.execute_reply.started":"2023-08-10T05:46:26.905547Z","shell.execute_reply":"2023-08-10T05:46:26.984946Z"},"trusted":true},"execution_count":8,"outputs":[]},{"cell_type":"code","source":"# train_utils/distributed_utils.py\n\nfrom collections import defaultdict, deque\nimport datetime\nimport pickle\nimport time\nimport errno\nimport os\n\nimport torch\nimport torch.distributed as dist\n\n\nclass SmoothedValue(object):\n    \"\"\"Track a series of values and provide access to smoothed values over a\n    window or the global series average.\n    \"\"\"\n    def __init__(self, window_size=20, fmt=None):\n        if fmt is None:\n            fmt = \"{value:.4f} ({global_avg:.4f}) hist: {all_hist_data}\"\n        self.deque = deque(maxlen=window_size)  # deque简单理解成加强版list\n        self.total = 0.0\n        self.count = 0\n        self.fmt = fmt\n\n    def update(self, value, n=1):\n        self.deque.append(value)\n        self.count += n\n        self.total += value * n\n\n    def synchronize_between_processes(self):\n        \"\"\"\n        Warning: does not synchronize the deque!\n        \"\"\"\n        if not is_dist_avail_and_initialized():\n            return\n        t = torch.tensor([self.count, self.total], dtype=torch.float64, device=\"cuda\")\n        dist.barrier()\n        dist.all_reduce(t)\n        t = t.tolist()\n        self.count = int(t[0])\n        self.total = t[1]\n\n    @property\n    def median(self):  # @property 是装饰器，这里可简单理解为增加median属性(只读)\n        d = torch.tensor(list(self.deque))\n        return d.median().item()\n\n    @property\n    def avg(self):\n        d = torch.tensor(list(self.deque), dtype=torch.float32)\n        return d.mean().item()\n\n    @property\n    def global_avg(self):\n        return self.total / self.count\n\n    @property\n    def max(self):\n        return max(self.deque)\n\n    @property\n    def value(self):\n        return self.deque[-1]\n\n    @property\n    def all_hist_data(self):\n        delimiter = \" \"\n        hist_data = []\n        for data in self.deque:\n            hist_data.append(\"{:.2f}\".format(data))\n\n        return \"[\" + delimiter.join(hist_data) + \"]\"\n\n    def __str__(self):\n        return self.fmt.format(\n            median=self.median,\n            avg=self.avg,\n            global_avg=self.global_avg,\n            max=self.max,\n            value=self.value,\n            all_hist_data=self.all_hist_data)\n\n\ndef all_gather(data):\n    \"\"\"\n    Run all_gather on arbitrary picklable data (not necessarily tensors)\n    Args:\n        data: any picklable object\n    Returns:\n        list[data]: list of data gathered from each rank\n    \"\"\"\n    world_size = get_world_size()\n    if world_size == 1:\n        return [data]\n\n    # serialized to a Tensor\n    buffer = pickle.dumps(data)\n    storage = torch.ByteStorage.from_buffer(buffer)\n    tensor = torch.ByteTensor(storage).to(\"cuda\")\n\n    # obtain Tensor size of each rank\n    local_size = torch.tensor([tensor.numel()], device=\"cuda\")\n    size_list = [torch.tensor([0], device=\"cuda\") for _ in range(world_size)]\n    dist.all_gather(size_list, local_size)\n    size_list = [int(size.item()) for size in size_list]\n    max_size = max(size_list)\n\n    # receiving Tensor from all ranks\n    # we pad the tensor because torch all_gather does not support\n    # gathering tensors of different shapes\n    tensor_list = []\n    for _ in size_list:\n        tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device=\"cuda\"))\n    if local_size != max_size:\n        padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device=\"cuda\")\n        tensor = torch.cat((tensor, padding), dim=0)\n    dist.all_gather(tensor_list, tensor)\n\n    data_list = []\n    for size, tensor in zip(size_list, tensor_list):\n        buffer = tensor.cpu().numpy().tobytes()[:size]\n        data_list.append(pickle.loads(buffer))\n\n    return data_list\n\n\ndef reduce_dict(input_dict, average=True):\n    \"\"\"\n    Args:\n        input_dict (dict): all the values will be reduced\n        average (bool): whether to do average or sum\n    Reduce the values in the dictionary from all processes so that all processes\n    have the averaged results. Returns a dict with the same fields as\n    input_dict, after reduction.\n    \"\"\"\n    world_size = get_world_size()\n    if world_size < 2:  # 单GPU的情况\n        return input_dict\n    with torch.no_grad():  # 多GPU的情况\n        names = []\n        values = []\n        # sort the keys so that they are consistent across processes\n        for k in sorted(input_dict.keys()):\n            names.append(k)\n            values.append(input_dict[k])\n        values = torch.stack(values, dim=0)\n        dist.all_reduce(values)\n        if average:\n            values /= world_size\n\n        reduced_dict = {k: v for k, v in zip(names, values)}\n        return reduced_dict\n\n\nclass MetricLogger(object):\n    def __init__(self, delimiter=\"\\t\"):\n        self.meters = defaultdict(SmoothedValue)\n        self.delimiter = delimiter\n\n    def update(self, **kwargs):\n        for k, v in kwargs.items():\n            if isinstance(v, torch.Tensor):\n                v = v.item()\n            assert isinstance(v, (float, int))\n            self.meters[k].update(v)\n\n    def __getattr__(self, attr):\n        if attr in self.meters:\n            return self.meters[attr]\n        if attr in self.__dict__:\n            return self.__dict__[attr]\n        raise AttributeError(\"'{}' object has no attribute '{}'\".format(\n            type(self).__name__, attr))\n\n    def __str__(self):\n        loss_str = []\n        for name, meter in self.meters.items():\n            loss_str.append(\n                \"{}: {}\".format(name, str(meter))\n            )\n        return self.delimiter.join(loss_str)\n\n    def synchronize_between_processes(self):\n        for meter in self.meters.values():\n            meter.synchronize_between_processes()\n\n    def add_meter(self, name, meter):\n        self.meters[name] = meter\n\n    def log_every(self, iterable, print_freq, header=None):\n        i = 0\n        if not header:\n            header = \"\"\n        start_time = time.time()\n        end = time.time()\n        iter_time = SmoothedValue(fmt='{avg:.4f}')\n        data_time = SmoothedValue(fmt='{avg:.4f}')\n        space_fmt = \":\" + str(len(str(len(iterable)))) + \"d\"\n        if torch.cuda.is_available():\n            log_msg = self.delimiter.join([header,\n                                           '[{0' + space_fmt + '}/{1}]',\n                                           'eta: {eta}',\n                                           '{meters}',\n                                           'time: {time}',\n                                           'data: {data}',\n                                           'max mem: {memory:.0f}'])\n        else:\n            log_msg = self.delimiter.join([header,\n                                           '[{0' + space_fmt + '}/{1}]',\n                                           'eta: {eta}',\n                                           '{meters}',\n                                           'time: {time}',\n                                           'data: {data}'])\n        MB = 1024.0 * 1024.0\n        for obj in iterable:\n            data_time.update(time.time() - end)\n            yield obj\n            iter_time.update(time.time() - end)\n            if i % print_freq == 0 or i == len(iterable) - 1:\n                eta_second = iter_time.global_avg * (len(iterable) - i)\n                eta_string = str(datetime.timedelta(seconds=eta_second))\n                if torch.cuda.is_available():\n                    print(log_msg.format(i, len(iterable),\n                                         eta=eta_string,\n                                         meters=str(self),\n                                         time=str(iter_time),\n                                         data=str(data_time),\n                                         memory=torch.cuda.max_memory_allocated() / MB))\n                else:\n                    print(log_msg.format(i, len(iterable),\n                                         eta=eta_string,\n                                         meters=str(self),\n                                         time=str(iter_time),\n                                         data=str(data_time)))\n            i += 1\n            end = time.time()\n        total_time = time.time() - start_time\n        total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n        print('{} Total time: {} ({:.4f} s / it)'.format(header,\n                                                         total_time_str,\n\n                                                         total_time / len(iterable)))\n\n\ndef warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor):\n\n    def f(x):\n        \"\"\"根据step数返回一个学习率倍率因子\"\"\"\n        if x >= warmup_iters:  # 当迭代数大于给定的warmup_iters时，倍率因子为1\n            return 1\n        alpha = float(x) / warmup_iters\n        # 迭代过程中倍率因子从warmup_factor -> 1\n        return warmup_factor * (1 - alpha) + alpha\n\n    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=f)\n\n\ndef mkdir(path):\n    try:\n        os.makedirs(path)\n    except OSError as e:\n        if e.errno != errno.EEXIST:\n            raise\n\n\ndef setup_for_distributed(is_master):\n    \"\"\"\n    This function disables when not in master process\n    \"\"\"\n    import builtins as __builtin__\n    builtin_print = __builtin__.print\n\n    def print(*args, **kwargs):\n        force = kwargs.pop('force', False)\n        if is_master or force:\n            builtin_print(*args, **kwargs)\n\n    __builtin__.print = print\n\n\ndef is_dist_avail_and_initialized():\n    \"\"\"检查是否支持分布式环境\"\"\"\n    if not dist.is_available():\n        return False\n    if not dist.is_initialized():\n        return False\n    return True\n\n\ndef get_world_size():\n    if not is_dist_avail_and_initialized():\n        return 1\n    return dist.get_world_size()\n\n\ndef get_rank():\n    if not is_dist_avail_and_initialized():\n        return 0\n    return dist.get_rank()\n\n\ndef is_main_process():\n    return get_rank() == 0\n\n\ndef save_on_master(*args, **kwargs):\n    if is_main_process():\n        torch.save(*args, **kwargs)\n\n\ndef init_distributed_mode(args):\n    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:\n        args.rank = int(os.environ[\"RANK\"])\n        args.world_size = int(os.environ['WORLD_SIZE'])\n        args.gpu = int(os.environ['LOCAL_RANK'])\n    elif 'SLURM_PROCID' in os.environ:\n        args.rank = int(os.environ['SLURM_PROCID'])\n        args.gpu = args.rank % torch.cuda.device_count()\n    else:\n        print('Not using distributed mode')\n        args.distributed = False\n        return\n\n    args.distributed = True\n\n    torch.cuda.set_device(args.gpu)\n    args.dist_backend = 'nccl'\n    print('| distributed init (rank {}): {}'.format(\n        args.rank, args.dist_url), flush=True)\n    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,\n                                         world_size=args.world_size, rank=args.rank)\n    torch.distributed.barrier()\n    setup_for_distributed(args.rank == 0)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:46:26.988403Z","iopub.execute_input":"2023-08-10T05:46:26.988850Z","iopub.status.idle":"2023-08-10T05:46:27.049527Z","shell.execute_reply.started":"2023-08-10T05:46:26.988814Z","shell.execute_reply":"2023-08-10T05:46:27.048588Z"},"trusted":true},"execution_count":9,"outputs":[]},{"cell_type":"code","source":"# train_utils/coco_eval.py\nimport json\nimport copy\nfrom collections import defaultdict\n\nimport numpy as np\nimport torch\n\nfrom pycocotools.cocoeval import COCOeval\nfrom pycocotools.coco import COCO\nimport pycocotools.mask as mask_util\n\n# from train_utils.distributed_utils import all_gather\n\n\nclass CocoEvaluator(object):\n    def __init__(self, coco_gt, iou_types):\n        assert isinstance(iou_types, (list, tuple))\n        coco_gt = copy.deepcopy(coco_gt)\n        self.coco_gt = coco_gt\n\n        self.iou_types = iou_types\n        self.coco_eval = {}\n        for iou_type in iou_types:\n            self.coco_eval[iou_type] = COCOeval(coco_gt, iouType=iou_type)\n\n        self.img_ids = []\n        self.eval_imgs = {k: [] for k in iou_types}\n\n    def update(self, predictions):\n        img_ids = list(np.unique(list(predictions.keys())))\n        self.img_ids.extend(img_ids)\n\n        for iou_type in self.iou_types:\n            results = self.prepare(predictions, iou_type)\n            coco_dt = loadRes(self.coco_gt, results) if results else COCO()\n            coco_eval = self.coco_eval[iou_type]\n\n            coco_eval.cocoDt = coco_dt\n            coco_eval.params.imgIds = list(img_ids)\n            img_ids, eval_imgs = evaluate_inner(coco_eval)\n\n            self.eval_imgs[iou_type].append(eval_imgs)\n\n    def synchronize_between_processes(self):\n        for iou_type in self.iou_types:\n            self.eval_imgs[iou_type] = np.concatenate(self.eval_imgs[iou_type], 2)\n            create_common_coco_eval(self.coco_eval[iou_type], self.img_ids, self.eval_imgs[iou_type])\n\n    def accumulate(self):\n        for coco_eval in self.coco_eval.values():\n            coco_eval.accumulate()\n\n    def summarize(self):\n        for iou_type, coco_eval in self.coco_eval.items():\n            print(\"IoU metric: {}\".format(iou_type))\n            coco_eval.summarize()\n\n    def prepare(self, predictions, iou_type):\n        if iou_type == \"bbox\":\n            return self.prepare_for_coco_detection(predictions)\n        elif iou_type == \"segm\":\n            return self.prepare_for_coco_segmentation(predictions)\n        elif iou_type == \"keypoints\":\n            return self.prepare_for_coco_keypoint(predictions)\n        else:\n            raise ValueError(\"Unknown iou type {}\".format(iou_type))\n\n    def prepare_for_coco_detection(self, predictions):\n        coco_results = []\n        for original_id, prediction in predictions.items():\n            if len(prediction) == 0:\n                continue\n\n            # xmin, ymin, xmax, ymax\n            boxes = prediction[\"boxes\"]\n            boxes = convert_to_xywh(boxes)\n            boxes = boxes.tolist()\n            scores = prediction[\"scores\"].tolist()\n            labels = prediction[\"labels\"].tolist()\n\n            coco_results.extend(\n                [\n                    {\n                        \"image_id\": original_id,\n                        \"category_id\": labels[k],\n                        \"bbox\": box,\n                        \"score\": scores[k],\n                    }\n                    for k, box in enumerate(boxes)\n                ]\n            )\n        return coco_results\n\n    def prepare_for_coco_segmentation(self, predictions):\n        coco_results = []\n        for original_id, prediction in predictions.items():\n            if len(prediction) == 0:\n                continue\n\n            scores = prediction[\"scores\"]\n            labels = prediction[\"labels\"]\n            masks = prediction[\"masks\"]\n\n            masks = masks > 0.5\n\n            scores = prediction[\"scores\"].tolist()\n            labels = prediction[\"labels\"].tolist()\n\n            rles = [\n                mask_util.encode(np.array(mask[0, :, :, np.newaxis], dtype=np.uint8, order=\"F\"))[0]\n                for mask in masks\n            ]\n            for rle in rles:\n                rle[\"counts\"] = rle[\"counts\"].decode(\"utf-8\")\n\n            coco_results.extend(\n                [\n                    {\n                        \"image_id\": original_id,\n                        \"category_id\": labels[k],\n                        \"segmentation\": rle,\n                        \"score\": scores[k],\n                    }\n                    for k, rle in enumerate(rles)\n                ]\n            )\n        return coco_results\n\n    def prepare_for_coco_keypoint(self, predictions):\n        coco_results = []\n        for original_id, prediction in predictions.items():\n            if len(prediction) == 0:\n                continue\n\n            boxes = prediction[\"boxes\"]\n            boxes = convert_to_xywh(boxes).tolist()\n            scores = prediction[\"scores\"].tolist()\n            labels = prediction[\"labels\"].tolist()\n            keypoints = prediction[\"keypoints\"]\n            keypoints = keypoints.flatten(start_dim=1).tolist()\n\n            coco_results.extend(\n                [\n                    {\n                        \"image_id\": original_id,\n                        \"category_id\": labels[k],\n                        'keypoints': keypoint,\n                        \"score\": scores[k],\n                    }\n                    for k, keypoint in enumerate(keypoints)\n                ]\n            )\n        return coco_results\n\n\ndef convert_to_xywh(boxes):\n    xmin, ymin, xmax, ymax = boxes.unbind(1)\n    return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=1)\n\n\ndef merge(img_ids, eval_imgs):\n    all_img_ids = all_gather(img_ids)\n    all_eval_imgs = all_gather(eval_imgs)\n\n    merged_img_ids = []\n    for p in all_img_ids:\n        merged_img_ids.extend(p)\n\n    merged_eval_imgs = []\n    for p in all_eval_imgs:\n        merged_eval_imgs.append(p)\n\n    merged_img_ids = np.array(merged_img_ids)\n    merged_eval_imgs = np.concatenate(merged_eval_imgs, 2)\n\n    # keep only unique (and in sorted order) images\n    merged_img_ids, idx = np.unique(merged_img_ids, return_index=True)\n    merged_eval_imgs = merged_eval_imgs[..., idx]\n\n    return merged_img_ids, merged_eval_imgs\n\n\ndef create_common_coco_eval(coco_eval, img_ids, eval_imgs):\n    img_ids, eval_imgs = merge(img_ids, eval_imgs)\n    img_ids = list(img_ids)\n    eval_imgs = list(eval_imgs.flatten())\n\n    coco_eval.evalImgs = eval_imgs\n    coco_eval.params.imgIds = img_ids\n    coco_eval._paramsEval = copy.deepcopy(coco_eval.params)\n\n\n#################################################################\n# From pycocotools, just removed the prints and fixed\n# a Python3 bug about unicode not defined\n#################################################################\n\n# Ideally, pycocotools wouldn't have hard-coded prints\n# so that we could avoid copy-pasting those two functions\n\ndef createIndex(self):\n    # create index\n    # print('creating index...')\n    anns, cats, imgs = {}, {}, {}\n    imgToAnns, catToImgs = defaultdict(list), defaultdict(list)\n    if 'annotations' in self.dataset:\n        for ann in self.dataset['annotations']:\n            imgToAnns[ann['image_id']].append(ann)\n            anns[ann['id']] = ann\n\n    if 'images' in self.dataset:\n        for img in self.dataset['images']:\n            imgs[img['id']] = img\n\n    if 'categories' in self.dataset:\n        for cat in self.dataset['categories']:\n            cats[cat['id']] = cat\n\n    if 'annotations' in self.dataset and 'categories' in self.dataset:\n        for ann in self.dataset['annotations']:\n            catToImgs[ann['category_id']].append(ann['image_id'])\n\n    # print('index created!')\n\n    # create class members\n    self.anns = anns\n    self.imgToAnns = imgToAnns\n    self.catToImgs = catToImgs\n    self.imgs = imgs\n    self.cats = cats\n\n\nmaskUtils = mask_util\n\n\ndef loadRes(self, resFile):\n    \"\"\"\n    Load result file and return a result api object.\n    :param   resFile (str)     : file name of result file\n    :return: res (obj)         : result api object\n    \"\"\"\n    res = COCO()\n    res.dataset['images'] = [img for img in self.dataset['images']]\n\n    # print('Loading and preparing results...')\n    # tic = time.time()\n    if isinstance(resFile, (str, bytes)):\n        anns = json.load(open(resFile))\n    elif type(resFile) == np.ndarray:\n        anns = self.loadNumpyAnnotations(resFile)\n    else:\n        anns = resFile\n    assert type(anns) == list, 'results in not an array of objects'\n    annsImgIds = [ann['image_id'] for ann in anns]\n    assert set(annsImgIds) == (set(annsImgIds) & set(self.getImgIds())), \\\n        'Results do not correspond to current coco set'\n    if 'caption' in anns[0]:\n        imgIds = set([img['id'] for img in res.dataset['images']]) & set([ann['image_id'] for ann in anns])\n        res.dataset['images'] = [img for img in res.dataset['images'] if img['id'] in imgIds]\n        for id, ann in enumerate(anns):\n            ann['id'] = id + 1\n    elif 'bbox' in anns[0] and not anns[0]['bbox'] == []:\n        res.dataset['categories'] = copy.deepcopy(self.dataset['categories'])\n        for id, ann in enumerate(anns):\n            bb = ann['bbox']\n            x1, x2, y1, y2 = [bb[0], bb[0] + bb[2], bb[1], bb[1] + bb[3]]\n            if 'segmentation' not in ann:\n                ann['segmentation'] = [[x1, y1, x1, y2, x2, y2, x2, y1]]\n            ann['area'] = bb[2] * bb[3]\n            ann['id'] = id + 1\n            ann['iscrowd'] = 0\n    elif 'segmentation' in anns[0]:\n        res.dataset['categories'] = copy.deepcopy(self.dataset['categories'])\n        for id, ann in enumerate(anns):\n            # now only support compressed RLE format as segmentation results\n            ann['area'] = maskUtils.area(ann['segmentation'])\n            if 'bbox' not in ann:\n                ann['bbox'] = maskUtils.toBbox(ann['segmentation'])\n            ann['id'] = id + 1\n            ann['iscrowd'] = 0\n    elif 'keypoints' in anns[0]:\n        res.dataset['categories'] = copy.deepcopy(self.dataset['categories'])\n        for id, ann in enumerate(anns):\n            s = ann['keypoints']\n            x = s[0::3]\n            y = s[1::3]\n            x1, x2, y1, y2 = np.min(x), np.max(x), np.min(y), np.max(y)\n            ann['area'] = (x2 - x1) * (y2 - y1)\n            ann['id'] = id + 1\n            ann['bbox'] = [x1, y1, x2 - x1, y2 - y1]\n    # print('DONE (t={:0.2f}s)'.format(time.time()- tic))\n\n    res.dataset['annotations'] = anns\n    createIndex(res)\n    return res\n\n\ndef evaluate_inner(self):\n    '''\n    Run per image evaluation on given images and store results (a list of dict) in self.evalImgs\n    :return: None\n    '''\n    # tic = time.time()\n    # print('Running per image evaluation...')\n    p = self.params\n    # add backward compatibility if useSegm is specified in params\n    if p.useSegm is not None:\n        p.iouType = 'segm' if p.useSegm == 1 else 'bbox'\n        print('useSegm (deprecated) is not None. Running {} evaluation'.format(p.iouType))\n    # print('Evaluate annotation type *{}*'.format(p.iouType))\n    p.imgIds = list(np.unique(p.imgIds))\n    if p.useCats:\n        p.catIds = list(np.unique(p.catIds))\n    p.maxDets = sorted(p.maxDets)\n    self.params = p\n\n    self._prepare()\n    # loop through images, area range, max detection number\n    catIds = p.catIds if p.useCats else [-1]\n\n    if p.iouType == 'segm' or p.iouType == 'bbox':\n        computeIoU = self.computeIoU\n    elif p.iouType == 'keypoints':\n        computeIoU = self.computeOks\n    self.ious = {\n        (imgId, catId): computeIoU(imgId, catId)\n        for imgId in p.imgIds\n        for catId in catIds}\n\n    evaluateImg = self.evaluateImg\n    maxDet = p.maxDets[-1]\n    evalImgs = [\n        evaluateImg(imgId, catId, areaRng, maxDet)\n        for catId in catIds\n        for areaRng in p.areaRng\n        for imgId in p.imgIds\n    ]\n    # this is NOT in the pycocotools code, but could be done outside\n    evalImgs = np.asarray(evalImgs).reshape(len(catIds), len(p.areaRng), len(p.imgIds))\n    self._paramsEval = copy.deepcopy(self.params)\n    # toc = time.time()\n    # print('DONE (t={:0.2f}s).'.format(toc-tic))\n    return p.imgIds, evalImgs\n\n#################################################################\n# end of straight copy from pycocotools, just removing the prints\n#################################################################","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:46:27.054680Z","iopub.execute_input":"2023-08-10T05:46:27.055502Z","iopub.status.idle":"2023-08-10T05:46:27.127114Z","shell.execute_reply.started":"2023-08-10T05:46:27.055444Z","shell.execute_reply":"2023-08-10T05:46:27.125979Z"},"trusted":true},"execution_count":10,"outputs":[]},{"cell_type":"code","source":"# train_utils/coco_utils.py\n\nfrom tqdm import tqdm\n\nimport torch\nimport torchvision\nimport torch.utils.data\nfrom pycocotools.coco import COCO\n\n\ndef convert_to_coco_api(ds):\n    coco_ds = COCO()\n    # annotation IDs need to start at 1, not 0\n    ann_id = 1\n    dataset = {'images': [], 'categories': [], 'annotations': []}\n    categories = set()\n    for img_idx in range(len(ds)):\n        # find better way to get target\n        targets = ds.coco_index(img_idx)\n        if not targets:\n            continue\n        image_id = targets[TARGET_FIELD_IMAGE_ID].item()\n        img_dict = {}\n        img_dict['id'] = image_id\n        img_dict['height'] = targets[TARGET_FIELD_HEIGHT_WIDTH][0]\n        img_dict['width'] = targets[TARGET_FIELD_HEIGHT_WIDTH][1]\n        dataset['images'].append(img_dict)\n\n        # xmin, ymin, xmax, ymax\n        bboxes = targets[TARGET_FIELD_TASK1_ANCHORS]\n\n        # (xmin, ymin, xmax, ymax) to (xmin, ymin, w, h)\n        bboxes[:, 2:] -= bboxes[:, :2]\n        # 将box的相对坐标信息（0-1）转为绝对值坐标\n        bboxes[:, [0, 2]] = bboxes[:, [0, 2]] * img_dict[\"width\"]\n        bboxes[:, [1, 3]] = bboxes[:, [1, 3]] * img_dict[\"height\"]\n        bboxes = bboxes.tolist()\n        labels = targets[TARGET_FIELD_TASK1_LABELS].tolist()\n        # 注意这里的boxes area也要进行转换，否则导致(small, medium, large)计算错误\n        areas = (targets[TARGET_FIELD_TASK1_AREA] * img_dict[\"width\"] * img_dict[\"height\"]).tolist()\n        iscrowd = targets[TARGET_FIELD_TASK1_ISCROWD].tolist()\n        num_objs = len(bboxes)\n        for i in range(num_objs):\n            ann = {}\n            ann['image_id'] = image_id\n            ann['bbox'] = bboxes[i]\n            ann['category_id'] = labels[i]\n            categories.add(labels[i])\n            ann['area'] = areas[i]\n            ann['iscrowd'] = iscrowd[i]\n            ann['id'] = ann_id\n            dataset['annotations'].append(ann)\n            ann_id += 1\n    dataset['categories'] = [{'id': i} for i in sorted(categories)]\n    coco_ds.dataset = dataset\n    coco_ds.createIndex()\n    return coco_ds\n\n\ndef get_coco_api_from_dataset(dataset):\n    for _ in range(10):\n        if isinstance(dataset, torchvision.datasets.CocoDetection):\n            break\n        if isinstance(dataset, torch.utils.data.Subset):\n            dataset = dataset.dataset\n    if isinstance(dataset, torchvision.datasets.CocoDetection):\n        return dataset.coco\n    return convert_to_coco_api(dataset)\n","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:46:27.128745Z","iopub.execute_input":"2023-08-10T05:46:27.129410Z","iopub.status.idle":"2023-08-10T05:46:27.148977Z","shell.execute_reply.started":"2023-08-10T05:46:27.129373Z","shell.execute_reply":"2023-08-10T05:46:27.147772Z"},"trusted":true},"execution_count":11,"outputs":[]},{"cell_type":"code","source":"# group_by_aspect_ratio.py\n\nimport bisect\nfrom collections import defaultdict\nimport copy\nfrom itertools import repeat, chain\nimport math\nimport numpy as np\n\nimport torch\nimport torch.utils.data\nfrom torch.utils.data.sampler import BatchSampler, Sampler\nfrom torch.utils.model_zoo import tqdm\nimport torchvision\n\nfrom PIL import Image\n\n\ndef _repeat_to_at_least(iterable, n):\n    repeat_times = math.ceil(n / len(iterable))\n    repeated = chain.from_iterable(repeat(iterable, repeat_times))\n    return list(repeated)\n\n\nclass GroupedBatchSampler(BatchSampler):\n    \"\"\"\n    Wraps another sampler to yield a mini-batch of indices.\n    It enforces that the batch only contain elements from the same group.\n    It also tries to provide mini-batches which follows an ordering which is\n    as close as possible to the ordering from the original sampler.\n    Arguments:\n        sampler (Sampler): Base sampler.\n        group_ids (list[int]): If the sampler produces indices in range [0, N),\n            `group_ids` must be a list of `N` ints which contains the group id of each sample.\n            The group ids must be a continuous set of integers starting from\n            0, i.e. they must be in the range [0, num_groups).\n        batch_size (int): Size of mini-batch.\n    \"\"\"\n    def __init__(self, sampler, group_ids, batch_size):\n        if not isinstance(sampler, Sampler):\n            raise ValueError(\n                \"sampler should be an instance of \"\n                \"torch.utils.data.Sampler, but got sampler={}\".format(sampler)\n            )\n        self.sampler = sampler\n        self.group_ids = group_ids\n        self.batch_size = batch_size\n\n    def __iter__(self):\n        buffer_per_group = defaultdict(list)\n        samples_per_group = defaultdict(list)\n\n        num_batches = 0\n        for idx in self.sampler:\n            group_id = self.group_ids[idx]\n            buffer_per_group[group_id].append(idx)\n            samples_per_group[group_id].append(idx)\n            if len(buffer_per_group[group_id]) == self.batch_size:\n                yield buffer_per_group[group_id]\n                num_batches += 1\n                del buffer_per_group[group_id]\n            assert len(buffer_per_group[group_id]) < self.batch_size\n\n        # now we have run out of elements that satisfy\n        # the group criteria, let's return the remaining\n        # elements so that the size of the sampler is\n        # deterministic\n        expected_num_batches = len(self)\n        num_remaining = expected_num_batches - num_batches\n        if num_remaining > 0:\n            # for the remaining batches, take first the buffers with largest number\n            # of elements\n            for group_id, _ in sorted(buffer_per_group.items(),\n                                      key=lambda x: len(x[1]), reverse=True):\n                remaining = self.batch_size - len(buffer_per_group[group_id])\n                samples_from_group_id = _repeat_to_at_least(samples_per_group[group_id], remaining)\n                buffer_per_group[group_id].extend(samples_from_group_id[:remaining])\n                assert len(buffer_per_group[group_id]) == self.batch_size\n                yield buffer_per_group[group_id]\n                num_remaining -= 1\n                if num_remaining == 0:\n                    break\n        assert num_remaining == 0\n\n    def __len__(self):\n        return len(self.sampler) // self.batch_size\n\n\ndef _compute_aspect_ratios_slow(dataset, indices=None):\n    print(\"Your dataset doesn't support the fast path for \"\n          \"computing the aspect ratios, so will iterate over \"\n          \"the full dataset and load every image instead. \"\n          \"This might take some time...\")\n    if indices is None:\n        indices = range(len(dataset))\n\n    class SubsetSampler(Sampler):\n        def __init__(self, indices):\n            self.indices = indices\n\n        def __iter__(self):\n            return iter(self.indices)\n\n        def __len__(self):\n            return len(self.indices)\n\n    sampler = SubsetSampler(indices)\n    data_loader = torch.utils.data.DataLoader(\n        dataset, batch_size=1, sampler=sampler,\n        num_workers=14,  # you might want to increase it for faster processing\n        collate_fn=lambda x: x[0])\n    aspect_ratios = []\n    with tqdm(total=len(dataset)) as pbar:\n        for _i, (img, _) in enumerate(data_loader):\n            pbar.update(1)\n            height, width = img.shape[-2:]\n            aspect_ratio = float(width) / float(height)\n            aspect_ratios.append(aspect_ratio)\n    return aspect_ratios\n\n\ndef _compute_aspect_ratios_custom_dataset(dataset, indices=None):\n    if indices is None:\n        indices = range(len(dataset))\n    aspect_ratios = []\n    for i in indices:\n        height, width = dataset.get_height_and_width(i)\n        aspect_ratio = float(width) / float(height)\n        aspect_ratios.append(aspect_ratio)\n    return aspect_ratios\n\n\ndef _compute_aspect_ratios_coco_dataset(dataset, indices=None):\n    if indices is None:\n        indices = range(len(dataset))\n    aspect_ratios = []\n    for i in indices:\n        img_info = dataset.coco.imgs[dataset.ids[i]]\n        aspect_ratio = float(img_info[\"width\"]) / float(img_info[\"height\"])\n        aspect_ratios.append(aspect_ratio)\n    return aspect_ratios\n\n\ndef _compute_aspect_ratios_voc_dataset(dataset, indices=None):\n    if indices is None:\n        indices = range(len(dataset))\n    aspect_ratios = []\n    for i in indices:\n        # this doesn't load the data into memory, because PIL loads it lazily\n        width, height = Image.open(dataset.images[i]).size\n        aspect_ratio = float(width) / float(height)\n        aspect_ratios.append(aspect_ratio)\n    return aspect_ratios\n\n\ndef _compute_aspect_ratios_subset_dataset(dataset, indices=None):\n    if indices is None:\n        indices = range(len(dataset))\n\n    ds_indices = [dataset.indices[i] for i in indices]\n    return compute_aspect_ratios(dataset.dataset, ds_indices)\n\n\ndef compute_aspect_ratios(dataset, indices=None):\n    if hasattr(dataset, \"get_height_and_width\"):\n        return _compute_aspect_ratios_custom_dataset(dataset, indices)\n\n    if isinstance(dataset, torchvision.datasets.CocoDetection):\n        return _compute_aspect_ratios_coco_dataset(dataset, indices)\n\n    if isinstance(dataset, torchvision.datasets.VOCDetection):\n        return _compute_aspect_ratios_voc_dataset(dataset, indices)\n\n    if isinstance(dataset, torch.utils.data.Subset):\n        return _compute_aspect_ratios_subset_dataset(dataset, indices)\n\n    # slow path\n    return _compute_aspect_ratios_slow(dataset, indices)\n\n\ndef _quantize(x, bins):\n    bins = copy.deepcopy(bins)\n    bins = sorted(bins)\n    # bisect_right：寻找y元素按顺序应该排在bins中哪个元素的右边，返回的是索引\n    quantized = list(map(lambda y: bisect.bisect_right(bins, y), x))\n    return quantized\n\n\ndef create_aspect_ratio_groups(dataset, k=0):\n    # 计算所有数据集中的图片width/height比例\n    aspect_ratios = compute_aspect_ratios(dataset)\n    # 将[0.5, 2]区间划分成2*k+1等份\n    bins = (2 ** np.linspace(-1, 1, 2 * k + 1)).tolist() if k > 0 else [1.0]\n\n    # 统计所有图像比例在bins区间中的位置索引\n    groups = _quantize(aspect_ratios, bins)\n    # count number of elements per group\n    # 统计每个区间的频次\n    counts = np.unique(groups, return_counts=True)[1]\n    fbins = [0] + bins + [np.inf]\n    print(\"Using {} as bins for aspect ratio quantization\".format(fbins))\n    print(\"Count of instances per bin: {}\".format(counts))\n    return groups","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:46:27.150726Z","iopub.execute_input":"2023-08-10T05:46:27.151302Z","iopub.status.idle":"2023-08-10T05:46:27.189300Z","shell.execute_reply.started":"2023-08-10T05:46:27.151269Z","shell.execute_reply":"2023-08-10T05:46:27.188211Z"},"trusted":true},"execution_count":12,"outputs":[]},{"cell_type":"code","source":"# train_utils/train_eval_utils.py\n\nimport math\nimport sys\nimport time\n\nimport torch\nimport torchvision\nfrom torchvision.utils import save_image\n\n# from train_utils import get_coco_api_from_dataset, CocoEvaluator\n# import train_utils.distributed_utils as utils\n\n\ndef train_one_epoch(model, optimizer, data_loader, device, epoch,\n                    print_freq=20, warmup=False):\n    model.train()\n    metric_logger = MetricLogger(delimiter=\"  \")\n    metric_logger.add_meter('lr', SmoothedValue(window_size=1, fmt='{value:.20f}'))\n    header = 'Epoch: [{}]'.format(epoch)\n    lr_scheduler = None\n    if epoch == 0 and warmup is True:  # 当训练第一轮（epoch=0）时，启用warmup训练方式，可理解为热身训练\n        warmup_factor = 5.0 / 10000\n        warmup_iters = min(1000, len(data_loader) - 1)\n\n        lr_scheduler = warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor)\n\n    mloss = torch.zeros(1).to(device)  # mean losses\n    for i, [images, targets] in enumerate(metric_logger.log_every(data_loader, print_freq, header)):\n        # batch inputs information\n        images = torch.stack(images, dim=0)\n        \n        grid_image = torchvision.utils.make_grid(images, nrow=4, padding=5)\n        save_image(grid_image, \"image.png\")\n        print(\"saving grid image\")\n\n        boxes = []\n        labels = []\n        img_id = []\n        task2_label = []\n        for t in targets:\n            boxes.append(t[TARGET_FIELD_TASK1_ANCHORS])\n            labels.append(t[TARGET_FIELD_TASK1_LABELS])\n            img_id.append(t[TARGET_FIELD_IMAGE_ID])\n            task2_label.append(t[TARGET_FIELD_TASK2_LABEL])\n        targets = {\"boxes\": torch.stack(boxes, dim=0),\n                   \"labels\": torch.stack(labels, dim=0),\n                   \"image_id\": torch.as_tensor(img_id),\n                   \"task2_label\":torch.as_tensor(task2_label)}\n\n        images = images.to(device)\n\n        targets = {k: v.to(device) for k, v in targets.items()}\n        losses_dict = model(epoch, images, targets)\n        losses = losses_dict[\"total_losses\"]\n\n        # reduce losses over all GPUs for logging purpose\n        losses_dict_reduced = reduce_dict(losses_dict)\n        losses_reduce = losses_dict_reduced[\"total_losses\"]\n\n        loss_value = losses_reduce.detach()\n        # 记录训练损失\n        mloss = (mloss * i + loss_value) / (i + 1)  # update mean losses\n\n        if not math.isfinite(loss_value):  # 当计算的损失为无穷大时停止训练\n            print(\"Loss is {}, stopping training\".format(loss_value))\n            print(losses_dict_reduced)\n            sys.exit(1)\n\n        optimizer.zero_grad()\n        losses.backward()\n        optimizer.step()\n\n        if lr_scheduler is not None:  # 第一轮使用warmup训练方式\n            lr_scheduler.step()\n\n        # metric_logger.update(loss=losses, **loss_dict_reduced)\n        metric_logger.update(**losses_dict_reduced)\n        now_lr = optimizer.param_groups[0][\"lr\"]\n        metric_logger.update(lr=now_lr)\n\n    return mloss, now_lr\n\n@torch.no_grad()\ndef evaluate(epoch, model, data_loader, device, data_set=None):\n\n    cpu_device = torch.device(\"cpu\")\n    model.eval()\n    metric_logger = MetricLogger(delimiter=\"  \")\n    header = \"Test: \"\n\n    if data_set is None:\n        data_set = get_coco_api_from_dataset(data_loader.dataset)\n    iou_types = _get_iou_types(model)\n    coco_evaluator = CocoEvaluator(data_set, iou_types)\n\n    task2_pred = []\n    task2_target = []\n    for images, targets in metric_logger.log_every(data_loader, 100, header):\n        images = torch.stack(images, dim=0).to(device)\n\n        if device != torch.device(\"cpu\"):\n            torch.cuda.synchronize(device)\n\n        model_time = time.time()\n        #  list((bboxes_out, labels_out, scores_out), ...)\n        results = model(epoch, images, targets=None)\n        task1_results = results[PREDICT_RESULT_TASK1]\n        task2_results = results[PREDICT_RESULT_TASK2]\n        \n        task2_pred.extend(task2_results.argmax(dim=1).tolist())\n        task2_target.extend([t[TARGET_FIELD_TASK2_LABEL].item() for t in targets])\n        \n        model_time = time.time() - model_time\n\n        outputs = []\n        for index, (bboxes_out, labels_out, scores_out) in enumerate(task1_results):\n            # 将box的相对坐标信息（0-1）转为绝对值坐标(xmin, ymin, xmax, ymax)\n            height_width = targets[index][TARGET_FIELD_HEIGHT_WIDTH]\n            \n            # 还原回原图尺度\n            gt_boxes = targets[index][TARGET_FIELD_TASK1_ANCHORS]\n            \n            gt_boxes[:, [0, 2]] = gt_boxes[:, [0, 2]] * height_width[1]\n            gt_boxes[:, [1, 3]] = gt_boxes[:, [1, 3]] * height_width[0]\n            \n            targets[index][TARGET_FIELD_TASK1_ANCHORS] = gt_boxes\n            \n            bboxes_out[:, [0, 2]] = bboxes_out[:, [0, 2]] * height_width[1]\n            bboxes_out[:, [1, 3]] = bboxes_out[:, [1, 3]] * height_width[0]\n\n            info = {\"boxes\": bboxes_out.to(cpu_device),\n                    \"labels\": labels_out.to(cpu_device),\n                    \"scores\": scores_out.to(cpu_device)}\n            outputs.append(info)\n\n        res = {target[TARGET_FIELD_IMAGE_ID].item(): output for target, output in zip(targets, outputs)}\n    \n        choice = random.randint(0, images.shape[0] - 1)\n        \n        height_width = targets[choice][TARGET_FIELD_HEIGHT_WIDTH]\n        \n        result = task1_results[choice]\n        bboxes = result[0]\n        bboxes[:, 0], bboxes[:, 2] = bboxes[:, 0] / height_width[1], bboxes[:, 2] / height_width[1]\n        bboxes[:, 1], bboxes[:, 3] = bboxes[:, 1] / height_width[0], bboxes[:, 3] / height_width[0]\n        \n        labels = result[1]\n        \n        gt_boxes = targets[index][TARGET_FIELD_TASK1_ANCHORS]\n        gt_boxes[:, 0], gt_boxes[:, 2] = gt_boxes[:, 0] / height_width[1], gt_boxes[:, 2] / height_width[1]\n        gt_boxes[:, 1], gt_boxes[:, 3] = gt_boxes[:, 1] / height_width[0], gt_boxes[:, 3] / height_width[0]\n        \n        gt_labels = targets[index][TARGET_FIELD_TASK1_LABELS]\n        filter_cond = gt_labels > 0\n        print(f\"evaluate: gtboxes:{gt_boxes[filter_cond]}, gt_labels:{gt_labels[filter_cond]}\")\n        print(f\"evaluate: bboxes:{bboxes[:5].detach()}, labels:{labels[:5].detach()}\")\n        \n        if epoch % 5 == 0:\n            target_bboxes = targets[choice][TARGET_FIELD_TASK1_ANCHORS]\n            targets_label = targets[choice][TARGET_FIELD_TASK1_LABELS]\n            save_anchor_pic(image=images, path=\"./output/evaluate_pred.jpg\", bboxes=bboxes, \n                            bboxes_label = labels, height=600, width=600, \n                            targets=target_bboxes, targets_label=targets_label,\n                            choice=choice)\n        \n        evaluator_time = time.time()\n        coco_evaluator.update(res)\n        evaluator_time = time.time() - evaluator_time\n        metric_logger.update(model_time=model_time, evaluator_time=evaluator_time)\n\n    # gather the stats from all processes\n    metric_logger.synchronize_between_processes()\n    print(\"Averaged stats:\", metric_logger)\n    coco_evaluator.synchronize_between_processes()\n\n    # accumulate predictions from all images\n    coco_evaluator.accumulate()\n    coco_evaluator.summarize()\n\n    task2_pred = np.array(task2_pred)\n    task2_target = np.array(task2_target)\n    \n    print(f\"task2_pred:{task2_pred}\")\n    print(f\"task2_target:{task2_target}\")\n    print(f\"task2准确率：{np.array(task2_pred == task2_target).mean()}\")\n    \n    coco_info = coco_evaluator.coco_eval[iou_types[0]].stats.tolist()  # numpy to list\n\n    return coco_info\n\n\ndef _get_iou_types(model):\n    model_without_ddp = model\n    if isinstance(model, torch.nn.parallel.DistributedDataParallel):\n        model_without_ddp = model.module\n    iou_types = [\"bbox\"]\n    return iou_types","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:58:11.128015Z","iopub.execute_input":"2023-08-10T05:58:11.128493Z","iopub.status.idle":"2023-08-10T05:58:11.173309Z","shell.execute_reply.started":"2023-08-10T05:58:11.128452Z","shell.execute_reply":"2023-08-10T05:58:11.172045Z"},"trusted":true},"execution_count":36,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import Dataset\nimport os\nimport torch\nimport csv\nfrom collections import defaultdict\nfrom PIL import Image\n\n# 任务1标注文件字段索引\nTASK1_CSV_FIELD_ID_INDEX = 0\nTASK1_CSV_FIELD_FILENAME_INDEX = 1\nTASK1_CSV_FIELD_LABEL_INDEX = 2\nTASK1_CSV_FIELD_XMIN_INDEX = 3\nTASK1_CSV_FIELD_YMIN_INDEX = 4\nTASK1_CSV_FIELD_XMAX_INDEX = 5\nTASK1_CSV_FIELD_YMAX_INDEX = 6\nTASK1_CSV_FIELD_HEIGHT_INDEX = 7\nTASK1_CSV_FIELD_WIDTH_INDEX = 8\nTASK1_CSV_FIELD_CHANNEL_INDEX = 9\nTASK1_CSV_FIELD_TASK_TYPE_INDEX=10\n# 任务2标注文件字段索引\nTASK2_CSV_FIELD_FILENAME_INDEX = 0\nTASK2_CSV_FIELD_LABEL_INDEX = 1\n\nTARGET_FIELD_TASK1_ANCHORS = \"task1_anchors\"\nTARGET_FIELD_TASK1_LABELS = \"task1_labels\"\nTARGET_FIELD_TASK2_LABEL = \"task2_label\"\nTARGET_FIELD_IMAGE_ID = \"image_id\"\nTARGET_FIELD_TASK1_AREA = \"task1_area\"\nTARGET_FIELD_TASK1_ISCROWD = \"task1_iscrowd\"\nTARGET_FIELD_HEIGHT_WIDTH = \"height_width\"\nTARGET_FIELD_IMAGE_DATA = \"image_data\"\nTARGET_FIELD_TASK_TYPE=\"task_type\"\n\nTASK_TYPE_ALL=1\nTASK_TYPE_ONLY_TASK2=2\n\nPREDICT_RESULT_TASK1 = \"task1_result\"\nPREDICT_RESULT_TASK2 = \"task2_result\"\n\nINVALID_FILE_INDEX = \"-1\"\n\nIMAGE_SIZE = 300\n\nTASK2_LABEL_COUNT = 2\n\n# 数据预处理\nproject_root = \"/kaggle/working/conveyer_belt_detector/\"\ndataset_root = \"/kaggle/input/conveyer-belt-detect/dataset/\"\n\n# 汇总的标注文件\ntask1_train_detect_anno = \"/kaggle/input/anno-data-new/train_info_all_orig.csv\"\n# task1_train_detect_anno = \"/kaggle/input/anno-data-new/train_info_all.csv\"\n# task1_train_detect_anno = \"/kaggle/input/anno-data-new/train_info_all_base_trans.csv\"\ntask1_val_detect_anno = os.path.join(project_root, \"data/task1/val/val_infos.csv\")\ntask2_train_classes_anno = \"/kaggle/input/anno-data-new/classes_all_orig.txt\"\n# task2_train_classes_anno = \"/kaggle/input/anno-data-new/classes_all.txt\"\n# task2_train_classes_anno = \"/kaggle/input/anno-data-new/classes_all_base_trans.txt\"\ntask2_val_classes_anno = os.path.join(project_root, \"data/task2/val/classes.txt\")\n\npretrain_path=\"/kaggle/input/nvidia-ssdpyt-fp32-190826pt/nvidia_ssdpyt_fp32_190826.pt\"\n\ndef make_dir(path):\n    dir_name = os.path.dirname(path)\n    if not os.path.exists(dir_name):\n        print(\"创建目录\", dir_name)\n        os.makedirs(dir_name)\n\nmake_dir(path=project_root)\nmake_dir(path=task1_train_detect_anno)\nmake_dir(path=task1_val_detect_anno)\nmake_dir(path=task2_train_classes_anno)\nmake_dir(path=task2_val_classes_anno)\n\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\ndef transform_task2_label(label):\n    return 1 if label == 3 else label\n\n\ndef recover_task2_label(label):\n    return 3 if label == 1 else label\n\n\nclass MyDataset(Dataset):\n    def __init__(self, task1_file=\"train_info_all.csv\", task2_file=\"classes_all.txt\", transforms=None,\n                 val_data_flag=False):\n        super(MyDataset, self).__init__()\n        self.task1_file = task1_file\n        self.task2_file = task2_file\n        self.transforms = transforms\n\n        self.images = []\n        self.task1_anchors = []\n\n        self.task1_labels = []\n        self.task2_labels = []\n\n        self.height_widths = []\n        # 默认全流程, 后续根据标注文件进行修改\n        self.task_types = []\n\n        self.max_anchors_count_in_image = 0\n\n        self.val_data_flag = val_data_flag\n\n        self._read()\n\n    def _read(self):\n        def generate_empty_task2_labels_item():\n            return {}\n\n        task2_data = defaultdict(generate_empty_task2_labels_item)\n        with open(file=self.task2_file, mode=\"r\", encoding=\"utf8\") as f:\n            while True:\n                line = f.readline().strip()\n                if len(line) <= 0:\n                    break\n\n                line = line.split(\",\")\n\n                task2_data[line[TASK2_CSV_FIELD_FILENAME_INDEX]] = int(line[TASK2_CSV_FIELD_LABEL_INDEX])\n\n        with open(file=self.task1_file, mode=\"r\", encoding=\"utf8\") as f:\n            line = f.readline()\n            last_index = INVALID_FILE_INDEX\n            anchors_count_in_image = 0\n            while True:\n                line = f.readline().strip()\n                if len(line) <= 0:\n                    break\n\n                line = line.split(\",\")\n\n                if line[TASK1_CSV_FIELD_ID_INDEX] != last_index:\n                    if last_index != INVALID_FILE_INDEX:\n                        if len(line) < 11 or int(line[10]) == 1:\n                            self.images.append(last_file_name)\n                            self.task1_anchors.append(task1_anchors)\n                            self.task1_labels.append(task1_labels)\n                            task2_label = task2_data[line[TASK1_CSV_FIELD_FILENAME_INDEX]]\n                            self.task2_labels.append(transform_task2_label(label=task2_label))\n                            self.height_widths.append(\n                                [int(line[TASK1_CSV_FIELD_HEIGHT_INDEX]), int(line[TASK1_CSV_FIELD_WIDTH_INDEX])])\n\n                            if not self.val_data_flag:\n                                self.task_types.append(int(line[TASK1_CSV_FIELD_TASK_TYPE_INDEX]))\n                        #                         print(task1_anchors)\n                        #                         print(task1_labels)\n                        #                         print(task2_data[line[TASK1_CSV_FIELD_FILENAME_INDEX]])\n                        #                         print(self.height_widths)\n\n                    task1_anchors = []\n                    task1_labels = []\n\n                    if self.max_anchors_count_in_image < anchors_count_in_image:\n                        self.max_anchors_count_in_image = anchors_count_in_image\n\n                    anchors_count_in_image = 0\n\n                anchor_str = line[TASK1_CSV_FIELD_XMIN_INDEX:TASK1_CSV_FIELD_YMAX_INDEX + 1]\n                task1_anchors.append([float(ele) for ele in anchor_str])\n                task1_labels.append(int(line[TASK1_CSV_FIELD_LABEL_INDEX]))\n\n                anchors_count_in_image += 1\n\n                last_index = line[TASK1_CSV_FIELD_ID_INDEX]\n                last_file_name = line[TASK1_CSV_FIELD_FILENAME_INDEX]\n\n    def __getitem__(self, idx):\n        image_id = torch.tensor([idx])\n        # 图像数据\n        image = self.images[idx]\n\n        # BGR => RGB\n        image = Image.open(fp=image).convert('RGB')\n\n        # 真实框（标注结果）\n        # (N, 4), N为标注框个数\n        task1_anchors = []\n        for anchor in self.task1_anchors[idx]:\n            task1_anchors.append(torch.tensor(data=anchor).float())\n        #         print(task1_anchors)\n\n        # 任务一标签\n        # (N), N为标注框个数\n        task1_labels = []\n        for label in self.task1_labels[idx]:\n            # 标签是long类型\n            task1_labels.append(torch.tensor(data=label).long())\n\n        task1_anchors = torch.stack(tensors=task1_anchors, dim=0)\n        task1_labels = torch.as_tensor(data=task1_labels, dtype=torch.int64)\n\n        task2_label = torch.as_tensor(data=self.task2_labels[idx], dtype=torch.int64)\n        image_id = torch.tensor(data=[image_id], dtype=torch.int64)\n        height_width = torch.as_tensor(data=self.height_widths[idx], dtype=torch.int64)\n\n        target = {}\n        target[TARGET_FIELD_TASK1_ANCHORS] = task1_anchors.to(device=device)\n        target[TARGET_FIELD_TASK1_LABELS] = task1_labels.to(device=device)\n        target[TARGET_FIELD_TASK2_LABEL] = task2_label.to(device=device)\n        target[TARGET_FIELD_IMAGE_ID] = image_id.to(device=device)\n        target[TARGET_FIELD_HEIGHT_WIDTH] = height_width.to(device=device)\n\n        if not self.val_data_flag:\n            task_type = torch.as_tensor(self.task_types[idx], dtype=torch.int64)\n            target[TARGET_FIELD_TASK_TYPE] = task_type.to(device=device)\n\n        # 图像预处理\n        if self.transforms is not None:\n            image, target = self.transforms(image, target)\n\n        # 返回结果\n        return image, target\n\n    def __len__(self):\n        return len(self.images)\n\n    @staticmethod\n    def collate_fn(batch):\n        images, targets = tuple(zip(*batch))\n        return images, targets\n\n    def coco_index(self, idx):\n        \"\"\"\n        该方法是专门为 pycocotools 统计标签信息准备，不对图像和标签作任何处理\n        由于不用去读取图片，可大幅缩减统计时间\n\n        Args:\n            idx: 输入需要获取图像的索引\n        \"\"\"\n        boxes = []\n        labels = []\n\n        boxes.extend(self.task1_anchors[idx])\n        labels.extend(self.task1_labels[idx])\n        iscrowds = [0 for _ in range(len(boxes))]\n\n        # convert everything into a torch.Tensor\n        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n        labels = torch.as_tensor(labels, dtype=torch.int64)\n        iscrowds = torch.as_tensor(iscrowds, dtype=torch.int64)\n        height_width = torch.as_tensor(self.height_widths[idx], dtype=torch.int64)\n        image_id = torch.tensor([idx])\n        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n\n        target = {}\n        target[TARGET_FIELD_TASK1_ANCHORS] = boxes.to(device=device)\n        target[TARGET_FIELD_TASK1_LABELS] = labels.to(device=device)\n        target[TARGET_FIELD_IMAGE_ID] = image_id.to(device=device)\n        target[TARGET_FIELD_TASK1_AREA] = area.to(device=device)\n        target[TARGET_FIELD_TASK1_ISCROWD] = iscrowds.to(device=device)\n        target[TARGET_FIELD_HEIGHT_WIDTH] = height_width.to(device=device)\n\n        return target","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:58:11.292264Z","iopub.execute_input":"2023-08-10T05:58:11.292777Z","iopub.status.idle":"2023-08-10T05:58:11.342224Z","shell.execute_reply.started":"2023-08-10T05:58:11.292737Z","shell.execute_reply":"2023-08-10T05:58:11.341329Z"},"trusted":true},"execution_count":37,"outputs":[]},{"cell_type":"code","source":"# transforms.py\n\nimport random\n\nimport numpy as np\nimport torch\nimport torchvision.transforms as t\nfrom torchvision.transforms import functional as F_torchvision\n\n# from src import dboxes300_coco, calc_iou_tensor, Encoder\n\n\nclass Compose(object):\n    \"\"\"组合多个transform函数\"\"\"\n    def __init__(self, transforms):\n        self.transforms = transforms\n\n    def __call__(self, image, target=None):\n        for trans in self.transforms:\n            image, target = trans(image, target)\n        return image, target\n\n\nclass ToTensor(object):\n    \"\"\"将PIL图像转为Tensor\"\"\"\n    def __call__(self, image, target):\n        image = F_torchvision.to_tensor(image).contiguous()\n        return image, target\n    \nclass RandomHorizontalFlip(object):\n    \"\"\"随机水平翻转图像以及bboxes,该方法应放在ToTensor后\"\"\"\n    def __init__(self, prob=0.5):\n        self.prob = prob\n\n    def __call__(self, image, target):\n        if random.random() < self.prob:\n            # height, width = image.shape[-2:]\n            image = image.flip(-1)  # 水平翻转图片\n            bbox = target[TARGET_FIELD_TASK1_ANCHORS]\n            # bbox: xmin, ymin, xmax, ymax\n            # bbox[:, [0, 2]] = width - bbox[:, [2, 0]]  # 翻转对应bbox坐标信息\n            bbox[:, [0, 2]] = 1.0 - bbox[:, [2, 0]]  # 翻转对应bbox坐标信息\n            target[TARGET_FIELD_TASK1_ANCHORS] = bbox\n        return image, target\n\n\nclass SSDCropping(object):\n    \"\"\"\n    根据原文，对图像进行裁剪,该方法应放在ToTensor前\n    Cropping for SSD, according to original paper\n    Choose between following 3 conditions:\n    1. Preserve the original image\n    2. Random crop minimum IoU is among 0.1, 0.3, 0.5, 0.7, 0.9\n    3. Random crop\n    Reference to https://github.com/chauhan-utk/src.DomainAdaptation\n    \"\"\"\n    def __init__(self):\n        self.sample_options = (\n            # 不做裁剪\n            None,\n            # 最小和最大 IoU\n            (0.1, None),\n            (0.3, None),\n            (0.5, None),\n            (0.7, None),\n            (0.9, None),\n            # 不做限制\n            (None, None),\n        )\n        self.dboxes = dboxes300_coco()\n        \n    \n\n    def __call__(self, image, target):\n        # 死循环，确保一定会返回结果\n        while True:\n            mode = random.choice(self.sample_options)\n\n            if mode is None:  # 不做随机裁剪处理\n                return image, target\n\n            htot, wtot = target[TARGET_FIELD_HEIGHT_WIDTH]\n\n            min_iou, max_iou = mode\n            min_iou = float('-inf') if min_iou is None else min_iou\n            max_iou = float('+inf') if max_iou is None else max_iou\n            # Implementation use 5 iteration to find possible candidate\n            for _ in range(5):\n                # 0.3*0.3 approx. 0.1\n                w = random.uniform(0.3, 1.0)\n                h = random.uniform(0.3, 1.0)\n                if w/h < 0.5 or w/h > 2:  # 保证宽高比例在0.5-2之间\n                    continue\n#                 w = random.uniform(0.4, 1.0)\n#                 h = w\n\n                # left 0 ~ wtot - w, top 0 ~ htot - h\n                left = random.uniform(0, 1.0 - w)\n                top = random.uniform(0, 1.0 - h)\n\n                right = left + w\n                bottom = top + h\n\n                # boxes的坐标是在0-1之间的\n                bboxes = target[TARGET_FIELD_TASK1_ANCHORS]\n                new_box = torch.tensor([[left, top, right, bottom]]).to(device=device)\n                ious = calc_iou_tensor(bboxes, new_box)\n                # tailor all the bboxes and return\n                # all(): Returns True if all elements in the tensor are True, False otherwise.\n                if not ((ious > min_iou) & (ious < max_iou)).all():\n                    continue\n\n                # 计算所有目标框的中心点\n                xc = 0.5 * (bboxes[:, 0] + bboxes[:, 2])\n                yc = 0.5 * (bboxes[:, 1] + bboxes[:, 3])\n                # 查看哪些目标框的中心点没有在被截取的图像中\n                masks = (xc > left) & (xc < right) & (yc > top) & (yc < bottom)\n\n                # 如果所有的gt box的中心点都不在采样的patch中，则重新找\n                if not masks.any():\n                    continue\n\n                # 修改采样patch中的所有gt box的坐标（防止出现越界的情况）\n                bboxes[bboxes[:, 0] < left, 0] = left\n                bboxes[bboxes[:, 1] < top, 1] = top\n                bboxes[bboxes[:, 2] > right, 2] = right\n                bboxes[bboxes[:, 3] > bottom, 3] = bottom\n\n                # 虑除不在采样patch中的gt box\n                bboxes = bboxes[masks, :]\n\n                # 获取在采样patch中的gt box的标签\n                labels = target[TARGET_FIELD_TASK1_LABELS]\n                labels = labels[masks]\n\n                # 裁剪 patch\n                left_idx = int(left * wtot)\n                top_idx = int(top * htot)\n                right_idx = int(right * wtot)\n                bottom_idx = int(bottom * htot)\n                image = image.crop((left_idx, top_idx, right_idx, bottom_idx))\n                # 调整裁剪后的bboxes坐标信息\n\n                bboxes[:, 0] = (bboxes[:, 0] - left) / w\n                bboxes[:, 1] = (bboxes[:, 1] - top) / h\n                bboxes[:, 2] = (bboxes[:, 2] - left) / w\n                bboxes[:, 3] = (bboxes[:, 3] - top) / h\n                \n#                 image, bboxes = self.pad_img(image, bboxes)\n\n                # 更新crop后的gt box坐标信息以及标签信息\n                target[TARGET_FIELD_TASK1_ANCHORS] = bboxes\n                target[TARGET_FIELD_TASK1_LABELS] = labels\n                return image, target\n\n\nclass Resize(object):\n    \"\"\"对图像进行resize处理,该方法应放在ToTensor前\"\"\"\n    def __init__(self, size=(600, 600)):\n        self.resize = t.Resize(size)\n        \n    def __call__(self, image, target):\n        image = self.resize(image)\n        return image, target\n\n\nclass ColorJitter(object):\n    \"\"\"对图像颜色信息进行随机调整,该方法应放在ToTensor前\"\"\"\n    def __init__(self, brightness=0.125, contrast=0.5, saturation=0.5, hue=0.05):\n        self.trans = t.ColorJitter(brightness, contrast, saturation, hue)\n\n    def __call__(self, image, target):\n        image = self.trans(image)\n        return image, target\n\n\nclass Normalization(object):\n    \"\"\"对图像标准化处理,该方法应放在ToTensor后\"\"\"\n    def __init__(self, mean=None, std=None):\n        if mean is None:\n            mean = [0.3441988, 0.34242108, 0.3464927]\n        if std is None:\n            std = [0.19682558, 0.19890308, 0.2000567]\n        self.normalize = t.Normalize(mean=mean, std=std)\n\n    def __call__(self, image, target):\n        image = self.normalize(image)\n        return image, target\n\n\nclass AssignGTtoDefaultBox(object):\n    \"\"\" 将 DefaultBox 与 GT进行匹配 \"\"\"\n    def __init__(self):\n        self.default_box = dboxes300_coco()\n        self.encoder = Encoder(self.default_box)\n\n    def __call__(self, image, target):\n        boxes = target[TARGET_FIELD_TASK1_ANCHORS]\n        labels = target[TARGET_FIELD_TASK1_LABELS]\n        # bboxes_out (Tensor 8732 x 4), labels_out (Tensor 8732)\n        bboxes_out, labels_out = self.encoder.encode(boxes, labels)\n        target[TARGET_FIELD_TASK1_ANCHORS] = bboxes_out\n        target[TARGET_FIELD_TASK1_LABELS] = labels_out\n        return image, target","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:58:11.486298Z","iopub.execute_input":"2023-08-10T05:58:11.487022Z","iopub.status.idle":"2023-08-10T05:58:11.524254Z","shell.execute_reply.started":"2023-08-10T05:58:11.486984Z","shell.execute_reply":"2023-08-10T05:58:11.522868Z"},"trusted":true},"execution_count":38,"outputs":[]},{"cell_type":"code","source":"# plot_curv.py\n\nimport datetime\nimport matplotlib.pyplot as plt\n\n\ndef plot_loss_and_lr(train_loss, learning_rate):\n    try:\n        x = list(range(len(train_loss)))\n        fig, ax1 = plt.subplots(1, 1)\n        ax1.plot(x, train_loss, 'r', label='loss')\n        ax1.set_xlabel(\"epoch\")\n        ax1.set_ylabel(\"loss\")\n        ax1.set_title(\"Train Loss and lr\")\n        plt.legend(loc='best')\n\n        ax2 = ax1.twinx()\n        ax2.plot(x, learning_rate, label='lr')\n        ax2.set_ylabel(\"learning rate\")\n        ax2.set_xlim(0, len(train_loss))  # 设置横坐标整数间隔\n        plt.legend(loc='best')\n\n        handles1, labels1 = ax1.get_legend_handles_labels()\n        handles2, labels2 = ax2.get_legend_handles_labels()\n        plt.legend(handles1 + handles2, labels1 + labels2, loc='upper right')\n\n        fig.subplots_adjust(right=0.8)  # 防止出现保存图片显示不全的情况\n        fig.savefig('./loss_and_lr{}.png'.format(datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")))\n        plt.close()\n        print(\"successful save loss curve! \")\n    except Exception as e:\n        print(e)\n\n\ndef plot_map(mAP):\n    try:\n        x = list(range(len(mAP)))\n        plt.plot(x, mAP, label='mAp')\n        plt.xlabel('epoch')\n        plt.ylabel('mAP')\n        plt.title('Eval mAP')\n        plt.xlim(0, len(mAP))\n        plt.legend(loc='best')\n        plt.savefig('./mAP.png')\n        plt.close()\n        print(\"successful save mAP curve!\")\n    except Exception as e:\n        print(e)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:58:11.693809Z","iopub.execute_input":"2023-08-10T05:58:11.694264Z","iopub.status.idle":"2023-08-10T05:58:11.708057Z","shell.execute_reply.started":"2023-08-10T05:58:11.694230Z","shell.execute_reply":"2023-08-10T05:58:11.706779Z"},"trusted":true},"execution_count":39,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# train_ssd300.py\nimport os\nimport datetime\n\nimport torch\n\nfrom torch.utils.data import DataLoader\n\nfrom PIL import ImageFile\nImageFile.LOAD_TRUNCATED_IMAGES = True\n\ndef create_model(num_classes=2 + 1):\n    # 先构建一个 backbone\n    backbone = Backbone()\n\n    # 再构建一个SSD300\n    model = SSD300(backbone=backbone, num_classes=num_classes)\n\n#     pre_ssd_path=\"/kaggle/input/nvidia-ssdpyt-fp32-190826pt/nvidia_ssdpyt_fp32_190826.pt\"\n    pre_ssd_path=\"/kaggle/input/nvidia-ssdpyt-fp32-190826pt/pretrain-ssd300-54.pth\"\n    if not os.path.exists(pre_ssd_path):\n        raise FileNotFoundError(\"pretrain-ssd300-54.pth not find in {}\".format(pre_ssd_path))\n    pre_model_dict = torch.load(pre_ssd_path, map_location='cpu')\n    pre_weights_dict = pre_model_dict[\"model\"]\n\n    # 删除类别预测器权重，注意，回归预测器的权重可以重用，因为不涉及num_classes\n    del_conf_loc_dict = {}\n    for k, v in pre_weights_dict.items():\n        split_key = k.split(\".\")\n        if \"conf\" in split_key:\n            continue\n        del_conf_loc_dict.update({k: v})\n\n    missing_keys, unexpected_keys = model.load_state_dict(del_conf_loc_dict, strict=False)\n    if len(missing_keys) != 0 or len(unexpected_keys) != 0:\n        print(\"missing_keys: \", missing_keys)\n        print(\"unexpected_keys: \", unexpected_keys)\n\n    return model\n\ndef main(parser_data):\n\n    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\n    # 保存训练完成之后的权重\n    if not os.path.exists(\"save_weights\"):\n        os.mkdir(\"save_weights\")\n\n\n    # 定义一个结果文件\n    results_file = \"results{}.txt\".format(datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\"))\n\n    \"\"\"\n        1，数据读取后的处理工作\n            - 类型转换\n            - 数据增强\n    \"\"\"\n\n    data_transform = {\n        \"train\": Compose([\n#             SSDCropping(), # 图像切割\n                          Resize(),   # 统一大小\n#                           ColorJitter(), # 颜色抖动\n                          ToTensor(),  # 转张量\n#                           RandomHorizontalFlip(), # 水平翻转\n#                           Normalization(), # 标准化\n                          AssignGTtoDefaultBox()]), # 处理目标框和锚框\n\n        \"val\": Compose([Resize(),\n                        ToTensor(),\n                        Normalization()])\n  }\n\n    # 定义训练集\n    train_dataset = MyDataset(task1_file=task1_train_detect_anno,\n                              task2_file=task2_train_classes_anno,\n                              transforms=data_transform[\"train\"])\n\n    # 注意训练时，batch_size必须大于1\n    batch_size = parser_data.batch_size\n\n    assert batch_size > 1, \"batch size must be greater than 1\"\n\n    # 防止最后一个batch_size=1，如果最后一个batch_size=1就舍去\n    drop_last = True if len(train_dataset) % batch_size == 1 else False\n\n    # 数据预处理多少线程\n    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers\n    train_data_loader = DataLoader(train_dataset,\n                                   batch_size=batch_size,\n                                   shuffle=True,\n                                   # num_workers=nw,\n                                   collate_fn=train_dataset.collate_fn,\n                                   drop_last=drop_last)\n\n    # 定义验证集\n    val_dataset = MyDataset(task1_file=task1_val_detect_anno,\n                            task2_file=task2_val_classes_anno,\n                            transforms=data_transform[\"val\"],\n                            val_data_flag=True)\n    val_data_loader = torch.utils.data.DataLoader(val_dataset,\n                                                  batch_size=batch_size,\n                                                  shuffle=False,\n                                                  # num_workers=nw,\n                                                  collate_fn=train_dataset.collate_fn)\n    # 定义模型\n    model = create_model(num_classes=args.num_classes + 1)\n    model.to(device)\n\n    # 定义优化器\n    params = [p for p in model.parameters() if p.requires_grad]\n    optimizer = torch.optim.SGD(params=params, lr=0.001,\n                                momentum=0.9, weight_decay=0.0005)\n    # 学习率调度器\n    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer=optimizer,\n                                                   step_size=5,\n                                                   gamma=0.8)\n\n    # 如果指定了上次训练保存的权重文件地址，则接着上次结果接着训练\n    if parser_data.resume != \"\":\n        checkpoint = torch.load(parser_data.resume, map_location='cpu')\n        model.load_state_dict(checkpoint['model'])\n        optimizer.load_state_dict(checkpoint['optimizer'])\n        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n        parser_data.start_epoch = checkpoint['epoch'] + 1\n        print(\"the training process from epoch{}...\".format(parser_data.start_epoch))\n        \n    optimizer.param_groups[0][\"lr\"] = 0.000512\n\n    train_loss = []\n    learning_rate = []\n    val_map = []\n\n    # 提前加载验证集数据，以免每次验证时都要重新加载一次数据，节省时间\n    val_data = get_coco_api_from_dataset(val_data_loader.dataset)\n\n    for epoch in range(parser_data.start_epoch, parser_data.epochs):\n        mean_loss, lr = train_one_epoch(model=model, optimizer=optimizer,\n                                        data_loader=train_data_loader,\n                                        device=device, epoch=epoch,\n                                        print_freq=10)\n        train_loss.append(mean_loss.item())\n        learning_rate.append(lr)\n        \n        print(f\"Epoch:{epoch} finish, mean_loss:{mean_loss}, lr:{lr}\")\n\n        # 更新学习率\n        lr_scheduler.step()\n\n        # 测试数据\n        coco_info = evaluate(epoch=epoch, model=model, data_loader=val_data_loader,\n                             device=device, data_set=val_data)\n\n        # write into txt\n        with open(results_file, \"a\") as f:\n            # 写入的数据包括coco指标还有loss和learning rate\n            result_info = [str(round(i, 4)) for i in coco_info + [mean_loss.item()]] + [str(round(lr, 6))]\n            txt = \"epoch:{} {}\".format(epoch, '  '.join(result_info))\n            f.write(txt + \"\\n\")\n\n        val_map.append(coco_info[1])  # pascal mAP\n\n        # save weights\n        save_files = {\n            'model': model.state_dict(),\n            'optimizer': optimizer.state_dict(),\n            'lr_scheduler': lr_scheduler.state_dict(),\n            'epoch': epoch}\n        torch.save(save_files, \"./save_weights/ssd300-{}.pth\".format(epoch))\n\n    # plot loss and lr curve\n    if len(train_loss) != 0 and len(learning_rate) != 0:\n        from plot_curve import plot_loss_and_lr\n        plot_loss_and_lr(train_loss, learning_rate)\n\n    # plot mAP curve\n    if len(val_map) != 0:\n        from plot_curve import plot_map\n        plot_map(val_map)\n\nif __name__ == '__main__':\n\n    import argparse\n\n    parser = argparse.ArgumentParser(description=__doc__)\n\n    # 检测的目标类别个数，不包括背景(替换：自己的检测类别)\n    parser.add_argument('--num_classes', default=2, type=int, help='num_classes')\n    # 文件保存地址\n    parser.add_argument('--output-dir', default='./save_weights', help='path where to save')\n    # 若需要接着上次训练，则指定上次训练保存权重文件地址 ./save_weights/ssd300-136.pth ./save_weights/ssd300-11.pth /kaggle/input/ssd300-train/ssd300-49.pth\n    parser.add_argument('--resume', default='./save_weights/ssd300-120.pth', type=str, help='resume from checkpoint')\n    # 指定接着从哪个epoch数开始训练\n    parser.add_argument('--start_epoch', default=121, type=int, help='start epoch')\n    # 训练的总epoch数\n    parser.add_argument('--epochs', default=300, type=int, metavar='N',\n                        help='number of total epochs to run')\n    # 训练的batch size\n    parser.add_argument('--batch_size', default=16, type=int, metavar='N',\n                        help='batch size when training.')\n\n    args = parser.parse_args(args=[])\n\n    # 检查保存权重文件夹是否存在，不存在则创建\n    if not os.path.exists(args.output_dir):\n        os.makedirs(args.output_dir)\n\n    main(args)\n","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:58:12.089697Z","iopub.execute_input":"2023-08-10T05:58:12.090596Z","iopub.status.idle":"2023-08-10T05:58:24.302938Z","shell.execute_reply.started":"2023-08-10T05:58:12.090507Z","shell.execute_reply":"2023-08-10T05:58:24.300741Z"},"trusted":true},"execution_count":40,"outputs":[{"name":"stdout","text":"missing_keys:  ['belt_pos.weight', 'belt_pos.bias', 'conf.0.weight', 'conf.0.bias', 'conf.1.weight', 'conf.1.bias', 'conf.2.weight', 'conf.2.bias', 'conf.3.weight', 'conf.3.bias', 'conf.4.weight', 'conf.4.bias', 'conf.5.weight', 'conf.5.bias']\nunexpected_keys:  []\nthe training process from epoch121...\ncreating index...\nindex created!\nsaving grid image\n","output_type":"stream"},{"traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)","Cell \u001b[0;32mIn[40], line 208\u001b[0m\n\u001b[1;32m    205\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mexists(args\u001b[38;5;241m.\u001b[39moutput_dir):\n\u001b[1;32m    206\u001b[0m     os\u001b[38;5;241m.\u001b[39mmakedirs(args\u001b[38;5;241m.\u001b[39moutput_dir)\n\u001b[0;32m--> 208\u001b[0m \u001b[43mmain\u001b[49m\u001b[43m(\u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n","Cell \u001b[0;32mIn[40], line 138\u001b[0m, in \u001b[0;36mmain\u001b[0;34m(parser_data)\u001b[0m\n\u001b[1;32m    135\u001b[0m val_data \u001b[38;5;241m=\u001b[39m get_coco_api_from_dataset(val_data_loader\u001b[38;5;241m.\u001b[39mdataset)\n\u001b[1;32m    137\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m epoch \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(parser_data\u001b[38;5;241m.\u001b[39mstart_epoch, parser_data\u001b[38;5;241m.\u001b[39mepochs):\n\u001b[0;32m--> 138\u001b[0m     mean_loss, lr \u001b[38;5;241m=\u001b[39m \u001b[43mtrain_one_epoch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    139\u001b[0m \u001b[43m                                    \u001b[49m\u001b[43mdata_loader\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtrain_data_loader\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    140\u001b[0m \u001b[43m                                    \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mepoch\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mepoch\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    141\u001b[0m \u001b[43m                                    \u001b[49m\u001b[43mprint_freq\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    142\u001b[0m     train_loss\u001b[38;5;241m.\u001b[39mappend(mean_loss\u001b[38;5;241m.\u001b[39mitem())\n\u001b[1;32m    143\u001b[0m     learning_rate\u001b[38;5;241m.\u001b[39mappend(lr)\n","Cell \u001b[0;32mIn[36], line 54\u001b[0m, in \u001b[0;36mtrain_one_epoch\u001b[0;34m(model, optimizer, data_loader, device, epoch, print_freq, warmup)\u001b[0m\n\u001b[1;32m     51\u001b[0m images \u001b[38;5;241m=\u001b[39m images\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[1;32m     53\u001b[0m targets \u001b[38;5;241m=\u001b[39m {k: v\u001b[38;5;241m.\u001b[39mto(device) \u001b[38;5;28;01mfor\u001b[39;00m k, v \u001b[38;5;129;01min\u001b[39;00m targets\u001b[38;5;241m.\u001b[39mitems()}\n\u001b[0;32m---> 54\u001b[0m losses_dict \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mepoch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mimages\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtargets\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     55\u001b[0m losses \u001b[38;5;241m=\u001b[39m losses_dict[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtotal_losses\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m     57\u001b[0m \u001b[38;5;66;03m# reduce losses over all GPUs for logging purpose\u001b[39;00m\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n","Cell \u001b[0;32mIn[8], line 168\u001b[0m, in \u001b[0;36mSSD300.forward\u001b[0;34m(self, eopch, image, targets)\u001b[0m\n\u001b[1;32m    166\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, eopch, image, targets\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[1;32m    167\u001b[0m     \u001b[38;5;66;03m# [b, 3, 600, 600] -> [b, 1024, 76, 76]\u001b[39;00m\n\u001b[0;32m--> 168\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfeature_extractor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimage\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    170\u001b[0m     shortcuts \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m    171\u001b[0m     shortcuts\u001b[38;5;241m.\u001b[39mappend(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mshortcut_blocks[\u001b[38;5;241m0\u001b[39m](x))\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n","Cell \u001b[0;32mIn[8], line 42\u001b[0m, in \u001b[0;36mBackbone.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     41\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x):\n\u001b[0;32m---> 42\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfeature_extractor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     43\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m x\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/nn/modules/container.py:217\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    215\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m):\n\u001b[1;32m    216\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m:\n\u001b[0;32m--> 217\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mmodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    218\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28minput\u001b[39m\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/nn/modules/container.py:217\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    215\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m):\n\u001b[1;32m    216\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m:\n\u001b[0;32m--> 217\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mmodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    218\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28minput\u001b[39m\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torchvision/models/resnet.py:158\u001b[0m, in \u001b[0;36mBottleneck.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    155\u001b[0m out \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbn3(out)\n\u001b[1;32m    157\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdownsample \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 158\u001b[0m     identity \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdownsample\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    160\u001b[0m out \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m identity\n\u001b[1;32m    161\u001b[0m out \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrelu(out)\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/nn/modules/container.py:217\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    215\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m):\n\u001b[1;32m    216\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m:\n\u001b[0;32m--> 217\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mmodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    218\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28minput\u001b[39m\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/nn/modules/conv.py:463\u001b[0m, in \u001b[0;36mConv2d.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    462\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[0;32m--> 463\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_conv_forward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbias\u001b[49m\u001b[43m)\u001b[49m\n","File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/nn/modules/conv.py:459\u001b[0m, in \u001b[0;36mConv2d._conv_forward\u001b[0;34m(self, input, weight, bias)\u001b[0m\n\u001b[1;32m    455\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpadding_mode \u001b[38;5;241m!=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mzeros\u001b[39m\u001b[38;5;124m'\u001b[39m:\n\u001b[1;32m    456\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m F\u001b[38;5;241m.\u001b[39mconv2d(F\u001b[38;5;241m.\u001b[39mpad(\u001b[38;5;28minput\u001b[39m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reversed_padding_repeated_twice, mode\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpadding_mode),\n\u001b[1;32m    457\u001b[0m                     weight, bias, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstride,\n\u001b[1;32m    458\u001b[0m                     _pair(\u001b[38;5;241m0\u001b[39m), \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdilation, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgroups)\n\u001b[0;32m--> 459\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconv2d\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbias\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstride\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    460\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpadding\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdilation\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgroups\u001b[49m\u001b[43m)\u001b[49m\n","\u001b[0;31mKeyboardInterrupt\u001b[0m: "],"ename":"KeyboardInterrupt","evalue":"","output_type":"error"}]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\n# 导出预测文件\ndef export_predict_result(image_root, train_weights, task1_file, task2_file):\n    make_dir(path=task1_file)\n    make_dir(path=task2_file)\n    # 判断设备\n    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\n    # 构建模型\n    # 目标检测数 + 背景\n    num_classes = 2 + 1\n    model = create_model(num_classes=num_classes)\n\n    # 加载权重\n    model.load_state_dict(torch.load(train_weights, map_location='cpu')['model'])\n    model.to(device)\n\n    model.eval()\n    with torch.no_grad():\n        # initial model 到底是为什么？\n        init_img = torch.zeros((1, 3, 600, 600), device=device)\n        model(1, init_img)\n        \n        with open(file=task1_file, mode=\"w\", encoding=\"utf8\", newline='') as task1_csv_file:\n            task1_header = [\"filename\", \"label\", \"xmin\", \"ymin\", \"width\", \"height\", \"confidence\"]\n            task1_csv_writer = csv.DictWriter(f=task1_csv_file, fieldnames=task1_header)\n            task1_csv_writer.writeheader()\n\n            with open(file=task2_file, mode=\"w\", encoding=\"utf8\", newline='') as task2_csv_file:\n#                 task2_header = [\"filename\", \"label\"]\n#                 task2_csv_writer = csv.DictWriter(f=task2_csv_file, fieldnames=task2_header)\n#                 task2_csv_writer.writeheader()\n            \n                for image_file_name in os.listdir(image_root):\n                    # 读取图像\n                    original_img = Image.open(os.path.join(image_root, image_file_name))\n\n                    data_transform = Compose([Resize(),\n                                              ToTensor(),\n                                              Normalization()])\n                    img, _ = data_transform(original_img)\n                    # 改为批量预测\n                    img = torch.unsqueeze(img, dim=0)\n\n                    time_start = time_synchronized()\n                    results = model(1, img.to(device))\n                    time_end = time_synchronized()\n                    print(\"{}, inference+NMS time: {}\".format(image_file_name, time_end - time_start))\n                    \n                    predictions = results[PREDICT_RESULT_TASK1][0]    # bboxes_out, labels_out, scores_out\n                    predict_boxes = predictions[0].to(\"cpu\").numpy()\n                    if len(predict_boxes) > 0:\n                        predict_boxes[:, [0, 2]] = predict_boxes[:, [0, 2]] * original_img.size[0]\n                        predict_boxes[:, [1, 3]] = predict_boxes[:, [1, 3]] * original_img.size[1]\n\n                        predict_classes = predictions[1].to(\"cpu\").numpy()\n                        predict_scores = predictions[2].to(\"cpu\").numpy()\n                        \n                        for i in range(len(predict_boxes)):\n                            if predict_scores[i] < 0.10:\n                                continue\n                            \n                            task1_new_row = defaultdict()\n                            \n                            task1_new_row[\"filename\"] = image_file_name\n                            task1_new_row[\"label\"] = predict_classes[i]\n\n                            task1_new_row[\"xmin\"] = int(predict_boxes[i, 0])\n                            task1_new_row[\"ymin\"] = int(predict_boxes[i, 1])\n                            task1_new_row[\"width\"] = int(predict_boxes[i, 2] - predict_boxes[i, 0])\n                            task1_new_row[\"height\"] = int(predict_boxes[i, 3] - predict_boxes[i, 1])\n\n                            task1_new_row[\"confidence\"] = round(predict_scores[i], 2)\n                            \n                            task1_csv_writer.writerow(task1_new_row)\n                    \n#                     task2_result = results[PREDICT_RESULT_TASK2].argmax(dim=1)[0]\n                    \n#                     task2_new_row = defaultdict()\n#                     task2_new_row[\"filename\"] = image_file_name\n#                     task2_new_row[\"label\"] = recover_task2_label(task2_result.item())\n                    \n#                     task2_csv_writer.writerow(task2_new_row)\n                    \n        print(\"导出 \" + task1_file + \" 成功\")\n        print(\"导出 \" + task2_file + \" 成功\")\n\n# test_image_root = os.path.join(dataset_root, \"test\", \"images\")\n# export_predict_result(image_root=test_image_root, \n#                       train_weights=\"./save_weights/ssd300-20.pth\",\n#                       task1_file=\"submit/submit1.csv\",\n#                       task2_file=\"submit/submit2.csv\")","metadata":{"execution":{"iopub.status.busy":"2023-08-10T05:54:32.517540Z","iopub.status.idle":"2023-08-10T05:54:32.517993Z","shell.execute_reply.started":"2023-08-10T05:54:32.517789Z","shell.execute_reply":"2023-08-10T05:54:32.517809Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}